Add data collection limit (#5137)

A-Artemis · jpggvilaca · web-flow · commit 48d32498e203 · 2026-01-12T13:52:02.000Z
Co-authored-by: Joao Vilaca &lt;joao.vilaca@intel.com&gt;
diff --git a/application/backend/app/__init__.py b/application/backend/app/__init__.py
@@ -1,3 +1,2 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
-
diff --git a/application/backend/app/alembic/versions/2786b50eb5a4_schema.py b/application/backend/app/alembic/versions/2786b50eb5a4_schema.py
@@ -148,7 +148,7 @@ def upgrade() -> None:
         sa.Column("sink_id", sa.Text(), nullable=True),
         sa.Column("model_revision_id", sa.Text(), nullable=True),
         sa.Column("is_running", sa.Boolean(), nullable=False),
-        sa.Column("data_collection_policies", sa.JSON(), nullable=False),
+        sa.Column("data_collection", sa.JSON(), nullable=False),
         sa.Column("device", sa.String(length=50), nullable=False),
         sa.Column("created_at", sa.DateTime(), server_default=sa.text("(CURRENT_TIMESTAMP)"), nullable=False),
         sa.Column("updated_at", sa.DateTime(), server_default=sa.text("(CURRENT_TIMESTAMP)"), nullable=False),
diff --git a/application/backend/app/api/__init__.py b/application/backend/app/api/__init__.py
@@ -1,3 +1,2 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
-
diff --git a/application/backend/app/api/routers/__init__.py b/application/backend/app/api/routers/__init__.py
@@ -1,3 +1,2 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
-
diff --git a/application/backend/app/api/routers/pipelines.py b/application/backend/app/api/routers/pipelines.py
@@ -13,13 +13,13 @@
 from app.api.dependencies import get_pipeline_metrics_service, get_pipeline_service, get_system_service
 from app.api.schemas import PipelineMetricsView, PipelineView
 from app.api.validators import ProjectID
-from app.models import DataCollectionPolicyAdapter, PipelineStatus
+from app.models import DataCollectionConfig, DataCollectionPolicyAdapter, PipelineStatus
 from app.services import PipelineMetricsService, PipelineService, ResourceNotFoundError, SystemService
 
 router = APIRouter(prefix="/api/projects/{project_id}/pipeline", tags=["Pipelines"])
 
 UPDATE_PIPELINE_BODY_DESCRIPTION = """
-Partial pipeline configuration update. May contain any subset of fields including 'device', 'data_collection_policies', 
+Partial pipeline configuration update. May contain any subset of fields including 'device', 'data_collection', 
 'source_id', 'sink_id', or 'model_id'. Fields not included in the request will remain unchanged.
 """
 UPDATE_PIPELINE_BODY_EXAMPLES = {
@@ -38,23 +38,26 @@
             "sink_id": "c6787c06-964b-4097-8eca-238b8cf79fc9",
         },
     ),
-    "enable_data_collection_policies": Example(
-        summary="Enable data collection policies",
-        description="Change data collection policies of the pipeline to fixed rate",
+    "enable_data_collection": Example(
+        summary="Enable data collection with max size",
+        description="Configure data collection with max dataset size and policies",
         value={
-            "data_collection_policies": [
-                {
-                    "type": "fixed_rate",
-                    "enabled": "true",
-                    "rate": 0.1,
-                }
-            ]
+            "data_collection": {
+                "max_dataset_size": 500,
+                "policies": [
+                    {
+                        "type": "fixed_rate",
+                        "enabled": True,
+                        "rate": 0.1,
+                    }
+                ],
+            }
         },
     ),
-    "clean_data_collection_policies": Example(
-        summary="Clean data collection policies",
+    "disable_data_collection": Example(
+        summary="Disable data collection",
         description="Remove all data collection policies of the pipeline",
-        value={"data_collection_policies": []},
+        value={"data_collection": {"max_dataset_size": None, "policies": []}},
     ),
     "change_device": Example(
         summary="Change inference device",
@@ -119,11 +122,13 @@ def update_pipeline(
             )
 
     try:
-        if "data_collection_policies" in pipeline_config:
-            pipeline_config["data_collection_policies"] = [
-                DataCollectionPolicyAdapter.validate_python(policy)
-                for policy in pipeline_config["data_collection_policies"]
-            ]
+        if "data_collection" in pipeline_config:
+            data_collection = pipeline_config["data_collection"]
+            if "policies" in data_collection:
+                data_collection["policies"] = [
+                    DataCollectionPolicyAdapter.validate_python(policy) for policy in data_collection["policies"]
+                ]
+            pipeline_config["data_collection"] = DataCollectionConfig.model_validate(data_collection)
         updated = pipeline_service.update_pipeline(project_id, pipeline_config)
         return PipelineView.model_validate(updated, from_attributes=True)
     except ResourceNotFoundError as e:
diff --git a/application/backend/app/api/schemas/pipeline.py b/application/backend/app/api/schemas/pipeline.py
@@ -5,7 +5,7 @@
 
 from pydantic import BaseModel, Field
 
-from app.models import DataCollectionPolicy, ModelRevision, PipelineStatus
+from app.models import DataCollectionConfig, ModelRevision, PipelineStatus
 
 from .sink import SinkView
 from .source import SourceView
@@ -17,7 +17,7 @@ class PipelineView(BaseModel):
     sink: SinkView | None = None  # None if disconnected
     model_revision: ModelRevision | None = Field(default=None, serialization_alias="model")
     status: PipelineStatus = PipelineStatus.IDLE
-    data_collection_policies: list[DataCollectionPolicy] = Field(default_factory=list)
+    data_collection: DataCollectionConfig = Field(default_factory=DataCollectionConfig)
     device: str = Field(default="cpu", description="Inference device (e.g., 'cpu', 'xpu', 'cuda', 'xpu-2', 'cuda-1')")
 
     model_config = {
@@ -54,19 +54,22 @@ class PipelineView(BaseModel):
                 },
                 "status": "running",
                 "device": "cpu",
-                "data_collection_policies": [
-                    {
-                        "type": "fixed_rate",
-                        "enabled": "true",
-                        "rate": 0.02,
-                    },
-                    {
-                        "type": "confidence_threshold",
-                        "enabled": "true",
-                        "confidence_threshold": 0.2,
-                        "min_sampling_interval": 2.5,
-                    },
-                ],
+                "data_collection": {
+                    "max_dataset_size": 500,
+                    "policies": [
+                        {
+                            "type": "fixed_rate",
+                            "enabled": True,
+                            "rate": 0.02,
+                        },
+                        {
+                            "type": "confidence_threshold",
+                            "enabled": True,
+                            "confidence_threshold": 0.2,
+                            "min_sampling_interval": 2.5,
+                        },
+                    ],
+                },
             }
         }
     }
diff --git a/application/backend/app/db/schema.py b/application/backend/app/db/schema.py
@@ -47,7 +47,7 @@ class PipelineDB(Base):
     sink_id: Mapped[str | None] = mapped_column(Text, ForeignKey("sinks.id", ondelete="RESTRICT"))
     model_revision_id: Mapped[str | None] = mapped_column(Text, ForeignKey("model_revisions.id", ondelete="RESTRICT"))
     is_running: Mapped[bool] = mapped_column(Boolean, default=False)
-    data_collection_policies: Mapped[list] = mapped_column(JSON, nullable=False, default=list)
+    data_collection: Mapped[dict] = mapped_column(JSON, nullable=False, default=dict)
     device: Mapped[str] = mapped_column(String(50), nullable=False, default="cpu")
 
     sink = relationship("SinkDB", uselist=False, lazy="joined")
diff --git a/application/backend/app/db_seeder.py b/application/backend/app/db_seeder.py
@@ -7,6 +7,7 @@
 
 from app.db.schema import LabelDB, ModelRevisionDB, PipelineDB, ProjectDB, SinkDB, SourceDB
 from app.models import (
+    DataCollectionConfig,
     DisconnectedSinkConfig,
     DisconnectedSourceConfig,
     FixedRateDataCollectionPolicy,
@@ -139,7 +140,10 @@ def _create_pipeline_with_video_source(  # noqa: PLR0913
     pipeline = PipelineDB(
         project_id=project_id,
         sink_id=sink_id,
-        data_collection_policies=[FixedRateDataCollectionPolicy(rate=0.1).model_dump(mode="json")],
+        data_collection=DataCollectionConfig(
+            max_dataset_size=100,
+            policies=[FixedRateDataCollectionPolicy(rate=0.1)],
+        ).model_dump(mode="json"),
         is_running=project_id == "9d6af8e8-6017-4ebe-9126-33aae739c5fa",  # Running only for detection project
     )
 
diff --git a/application/backend/app/models/__init__.py b/application/backend/app/models/__init__.py
@@ -3,6 +3,7 @@
 
 from .data_collection_policy import (
     ConfidenceThresholdDataCollectionPolicy,
+    DataCollectionConfig,
     DataCollectionPolicy,
     DataCollectionPolicyAdapter,
     FixedRateDataCollectionPolicy,
@@ -44,6 +45,7 @@
 
 __all__ = [
     "ConfidenceThresholdDataCollectionPolicy",
+    "DataCollectionConfig",
     "DataCollectionPolicy",
     "DataCollectionPolicyAdapter",
     "DatasetItem",
diff --git a/application/backend/app/models/data_collection_policy.py b/application/backend/app/models/data_collection_policy.py
@@ -27,3 +27,22 @@ class ConfidenceThresholdDataCollectionPolicy(DataCollectionPolicyBase):
 ]
 
 DataCollectionPolicyAdapter: TypeAdapter[DataCollectionPolicy] = TypeAdapter(DataCollectionPolicy)
+
+
+class DataCollectionConfig(BaseModel):
+    """
+    Configuration for data collection during pipeline execution.
+
+    Attributes:
+        max_dataset_size: Maximum number of items allowed in the dataset. When reached,
+            data collection will be disabled to prevent uncontrolled dataset growth.
+            Set to None for unlimited collection (default).
+        policies: List of policies governing data collection behavior.
+    """
+
+    max_dataset_size: int | None = Field(
+        default=None,
+        ge=1,
+        description="Maximum number of items allowed in the dataset. None for unlimited.",
+    )
+    policies: list[DataCollectionPolicy] = Field(default_factory=list)
diff --git a/application/backend/app/models/pipeline.py b/application/backend/app/models/pipeline.py
@@ -8,7 +8,7 @@
 from pydantic import AliasChoices, Field, model_validator
 
 from .base import BaseEntity
-from .data_collection_policy import DataCollectionPolicy
+from .data_collection_policy import DataCollectionConfig, DataCollectionPolicy
 from .model_revision import ModelRevision
 from .sink import Sink
 from .source import Source
@@ -43,7 +43,7 @@ class Pipeline(BaseEntity):
         sink_id: UUID reference to the sink entity.
         model_id: UUID reference to the model revision entity.
         status: Current operational status of the pipeline (IDLE or RUNNING).
-        data_collection_policies: List of policies governing data collection behavior during pipeline execution.
+        data_collection: Configuration for data collection including max dataset size and policies.
         device: The device used for model inference (e.g., 'cpu', 'xpu', 'cuda', 'xpu-1', etc.).
 
     Raises:
@@ -58,9 +58,14 @@ class Pipeline(BaseEntity):
     sink_id: UUID | None = None
     model_id: UUID | None = Field(default=None, validation_alias=AliasChoices("model_revision_id", "model_id"))
     status: PipelineStatus = PipelineStatus.IDLE
-    data_collection_policies: list[DataCollectionPolicy] = Field(default_factory=list)
+    data_collection: DataCollectionConfig = Field(default_factory=DataCollectionConfig)
     device: str = Field(default="cpu", pattern=r"^(cpu|xpu|cuda)(-\d+)?$")
 
+    @property
+    def data_collection_policies(self) -> list[DataCollectionPolicy]:
+        """Backward-compatible property to access data collection policies."""
+        return self.data_collection.policies
+
     @model_validator(mode="before")
     def set_status_from_is_running(cls, data: Any) -> Any:
         if hasattr(data, "is_running") and not hasattr(data, "status"):
diff --git a/application/backend/app/services/data_collect/data_collector.py b/application/backend/app/services/data_collect/data_collector.py
@@ -124,12 +124,12 @@ def _load_pipeline(self) -> None:
 
             self.active_pipeline_data = pipeline, project
             logger.info(
-                "Dataset collection policies set to {}, source: {}",
-                pipeline.data_collection_policies,
+                "Data collection config set to {}, source: {}",
+                pipeline.data_collection,
                 pipeline.source_id,
             )
 
-            policies = [policy for policy in pipeline.data_collection_policies if policy.enabled]
+            policies = [policy for policy in pipeline.data_collection.policies if policy.enabled]
             self.policy_checkers = []
             for policy in policies:
                 checker: PolicyChecker | None = None
@@ -156,7 +156,6 @@ def collect(
 
         Args:
             timestamp: Floating-point timestamp of the captured image, used for item naming.
-            confidence: Floating-point confidence of the captured image, used for item naming.
             frame_data: Image data in numpy ndarray format (expected in BGR color space).
             inference_data: Inference data containing model predictions and model identifier.
 
@@ -167,6 +166,7 @@ def collect(
             Collection occurs if any policy checker returns True OR if the
             should_collect_next_frame flag is set. Timestamp is formatted to string
             with 4 decimal places for use as dataset item name.
+            Collection is skipped if max_dataset_size is set and the dataset has reached that limit.
         """
         if self.active_pipeline_data is None:
             return
@@ -186,10 +186,23 @@ def collect(
         frame_data = cv2.cvtColor(frame_data, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
         with get_db_session() as session:
             label_service = LabelService(db_session=session)
+            dataset_service = DatasetService(data_dir=self.data_dir, label_service=label_service, db_session=session)
+
+            # Check if max_dataset_size limit has been reached
+            max_dataset_size = pipeline.data_collection.max_dataset_size
+            if max_dataset_size is not None:
+                current_count = dataset_service.count_dataset_items(project=project)
+                if current_count >= max_dataset_size:
+                    logger.debug(
+                        "Dataset has reached max size limit ({}/{}), skipping data collection",
+                        current_count,
+                        max_dataset_size,
+                    )
+                    return
+
             labels = label_service.list_all(project_id=project.id)
             annotations = convert_prediction(labels=labels, frame_data=frame_data, prediction=inference_data.prediction)
 
-            dataset_service = DatasetService(data_dir=self.data_dir, label_service=label_service, db_session=session)
             dataset_service.create_dataset_item(
                 project=project,
                 name=f"{timestamp:.4f}".replace(".", "_"),
diff --git a/application/backend/app/services/event/__init__.py b/application/backend/app/services/event/__init__.py
@@ -1,3 +1,2 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
-
diff --git a/application/backend/app/services/pipeline_service.py b/application/backend/app/services/pipeline_service.py
@@ -75,7 +75,7 @@ def update_pipeline(self, project_id: UUID, partial_config: dict) -> Pipeline:
             sink_id=str(to_update.sink_id) if to_update.sink_id else None,
             model_revision_id=str(to_update.model_id) if to_update.model_id else None,
             is_running=to_update.status.as_bool,
-            data_collection_policies=[obj.model_dump() for obj in to_update.data_collection_policies],
+            data_collection=to_update.data_collection.model_dump(),
             device=to_update.device,
         )
         pipeline_db = pipeline_repo.update(to_update_db)
@@ -86,7 +86,7 @@ def update_pipeline(self, project_id: UUID, partial_config: dict) -> Pipeline:
                 self._event_bus.emit_event(EventType.SOURCE_CHANGED)
             if pipeline.sink_id != updated.sink_id:  # type: ignore[union-attr] # sink is always there for running pipeline
                 self._event_bus.emit_event(EventType.SINK_CHANGED)
-            if pipeline.data_collection_policies != updated.data_collection_policies:
+            if pipeline.data_collection != updated.data_collection:
                 self._event_bus.emit_event(EventType.PIPELINE_DATASET_COLLECTION_POLICIES_CHANGED)
             if pipeline.device != updated.device:
                 self._event_bus.emit_event(EventType.INFERENCE_DEVICE_CHANGED)
diff --git a/application/backend/app/stream/__init__.py b/application/backend/app/stream/__init__.py
@@ -1,3 +1,2 @@
 # Copyright (C) 2025 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
-
diff --git a/application/backend/tests/integration/project_factory.py b/application/backend/tests/integration/project_factory.py
@@ -143,7 +143,7 @@ def with_data_policies(self, data_policies: list[dict]) -> "ProjectTestDataFacto
         """Set data collection policy for the project."""
         if not self._pipeline:
             raise ValueError("Pipeline must be set before adding data policies")
-        self._pipeline.data_collection_policies = data_policies
+        self._pipeline.data_collection = {"max_dataset_size": None, "policies": data_policies}
         return self
 
     def build(self) -> ProjectDB:
diff --git a/application/backend/tests/integration/services/test_pipeline_service.py b/application/backend/tests/integration/services/test_pipeline_service.py
diff --git a/application/backend/tests/unit/routers/test_pipelines.py b/application/backend/tests/unit/routers/test_pipelines.py
diff --git a/application/backend/tests/unit/services/data_collect/test_data_collector.py b/application/backend/tests/unit/services/data_collect/test_data_collector.py
diff --git a/application/ui/mocks/mock-pipeline.ts b/application/ui/mocks/mock-pipeline.ts
diff --git a/application/ui/src/features/inference/aside/data-collection.component.tsx b/application/ui/src/features/inference/aside/data-collection.component.tsx
diff --git a/application/ui/src/features/project/details/project-details.component.test.tsx b/application/ui/src/features/project/details/project-details.component.test.tsx
diff --git a/application/ui/tests/fixtures.ts b/application/ui/tests/fixtures.ts
diff --git a/application/ui/tests/inference/inference.spec.ts b/application/ui/tests/inference/inference.spec.ts

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,2 @@`
`1`	`1`	`# Copyright (C) 2025 Intel Corporation`
`2`	`2`	`# SPDX-License-Identifier: Apache-2.0`
`3`		`-`