fix: downgrade processing graph and end_time to optional, add auto-so… (#1529)

* fix: downgrade processing graph and end_time to optional, add auto-sorting

* chore: lint

* fix: handle empty data_prcesses

* tests: remove check for required field that is now optional

* tests: coverage on the new ordering validator

Author: dbirman

src/aind_data_schema/core/processing.py

@@ -65,7 +65,9 @@ class DataProcess(DataModel):
         default=None, title="Pipeline name", description="Pipeline names must exist in Processing.pipelines"
     )
     start_date_time: Annotated[AwareDatetimeWithDefault, TimeValidation.AFTER] = Field(..., title="Start date time")
-    end_date_time: Annotated[AwareDatetimeWithDefault, TimeValidation.AFTER] = Field(..., title="End date time")
+    end_date_time: Optional[Annotated[AwareDatetimeWithDefault, TimeValidation.AFTER]] = Field(
+        default=None, title="End date time"
+    )
     output_path: Optional[AssetPath] = Field(
         default=None, title="Output path", description="Path to processing outputs, if stored."
     )
@@ -110,8 +112,8 @@ class Processing(DataCoreModel):
     )
     notes: Optional[str] = Field(default=None, title="Notes")
 
-    dependency_graph: Dict[str, List[str]] = Field(
-        ...,
+    dependency_graph: Optional[Dict[str, List[str]]] = Field(
+        default=None,
         title="Dependency graph",
         description=(
             "Directed graph of processing step dependencies. Each key is a process name, and the value is a list of "
@@ -140,6 +142,26 @@ def rename_process(self, old_name: str, new_name: str) -> None:
             if old_name in value:
                 value[value.index(old_name)] = new_name
 
+    @model_validator(mode="after")
+    def order_processes(self) -> "Processing":
+        """Ensure that processes are ordered by start_date_time"""
+
+        if not hasattr(self, "data_processes") or not self.data_processes:
+            return self
+
+        # Check if any processes are out of order
+        start_times = [process.start_date_time for process in self.data_processes]
+        if not all(start_times[i] <= start_times[i + 1] for i in range(len(start_times) - 1)):
+            # Sort processes by start_date_time
+            self.data_processes.sort(key=lambda x: x.start_date_time)
+            self.notes = (
+                "Processes were reordered by start_date_time"
+                if not self.notes
+                else f"{self.notes}; Processes were reordered by start_date_time"
+            )
+
+        return self
+
     @classmethod
     def create_with_sequential_process_graph(cls, data_processes: List[DataProcess], **kwargs) -> "Processing":
         """Generate a sequential process graph from a list of DataProcess objects"""

tests/test_processing.py

@@ -60,7 +60,6 @@ def test_constructors(self):
         self.assertIn("stage", repr(e.exception))
         self.assertIn("code", repr(e.exception))
         self.assertIn("start_date_time", repr(e.exception))
-        self.assertIn("end_date_time", repr(e.exception))
         self.assertIn("notes", repr(e.exception))
 
     def test_resource_usage(self):
@@ -361,6 +360,96 @@ def test_validate_pipeline_names(self):
             )
         self.assertIn("Pipeline name 'NonExistentPipeline' not found in pipelines list", str(e.exception))
 
+    def test_order_processes(self):
+        """Test the order_processes method"""
+
+        # Create processes with different start times (out of order)
+        t1 = datetime.fromisoformat("2024-09-13T14:00:00")
+        t2 = datetime.fromisoformat("2024-09-13T12:00:00")  # Earlier time
+        t3 = datetime.fromisoformat("2024-09-13T16:00:00")  # Later time
+
+        process1 = DataProcess(
+            name="process1",
+            experimenters=["Dr. Dan"],
+            process_type=ProcessName.COMPRESSION,
+            stage=ProcessStage.PROCESSING,
+            code=code,
+            start_date_time=t1,
+            end_date_time=t1,
+        )
+        process2 = DataProcess(
+            name="process2",
+            experimenters=["Dr. Dan"],
+            process_type=ProcessName.ANALYSIS,
+            stage=ProcessStage.ANALYSIS,
+            code=code,
+            start_date_time=t2,
+            end_date_time=t2,
+        )
+        process3 = DataProcess(
+            name="process3",
+            experimenters=["Dr. Dan"],
+            process_type=ProcessName.SPIKE_SORTING,
+            stage=ProcessStage.PROCESSING,
+            code=code,
+            start_date_time=t3,
+            end_date_time=t3,
+        )
+
+        # Create Processing with out-of-order processes
+        dependency_graph = {"process1": [], "process2": [], "process3": []}
+        p = Processing(data_processes=[process1, process2, process3], dependency_graph=dependency_graph)
+
+        # Check that processes were reordered by start_date_time
+        expected_order = [process2, process1, process3]  # t2 < t1 < t3
+        actual_names = [proc.name for proc in p.data_processes]
+        expected_names = [proc.name for proc in expected_order]
+        self.assertEqual(actual_names, expected_names)
+
+        # Check that notes were updated
+        self.assertIn("Processes were reordered by start_date_time", p.notes)
+
+        # Test with already ordered processes
+        process4 = DataProcess(
+            name="process4",
+            experimenters=["Dr. Dan"],
+            process_type=ProcessName.COMPRESSION,
+            stage=ProcessStage.PROCESSING,
+            code=code,
+            start_date_time=t1,
+            end_date_time=t1,
+        )
+        process5 = DataProcess(
+            name="process5",
+            experimenters=["Dr. Dan"],
+            process_type=ProcessName.ANALYSIS,
+            stage=ProcessStage.ANALYSIS,
+            code=code,
+            start_date_time=t3,
+            end_date_time=t3,
+        )
+
+        dependency_graph2 = {"process4": [], "process5": []}
+        p2 = Processing(data_processes=[process4, process5], dependency_graph=dependency_graph2)
+
+        # Check that order wasn't changed and no reordering note was added
+        self.assertEqual([proc.name for proc in p2.data_processes], ["process4", "process5"])
+        self.assertIsNone(p2.notes)
+
+        # Test with existing notes
+        dependency_graph3 = {"process1": [], "process2": [], "process3": []}
+        p3 = Processing(
+            data_processes=[process1, process2, process3], dependency_graph=dependency_graph3, notes="Existing notes"
+        )
+
+        # Check that reordering note was appended to existing notes
+        self.assertIn("Existing notes; Processes were reordered by start_date_time", p3.notes)
+
+        # Test with empty data_processes
+        p4 = Processing(data_processes=[], dependency_graph={})
+        self.assertEqual(len(p4.data_processes), 0)
+        self.assertIsNone(p4.notes)
+
 
 if __name__ == "__main__":
     unittest.main()