fix: sort remaining metadata in topo sort

Author: Michał Sośnicki

src/neptune_exporter/loader_manager.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import datetime
 import heapq
 from pathlib import Path
 from typing import NewType, Optional
@@ -123,9 +124,9 @@ def _topological_sort_runs(
             # Otherwise, run is a root (orphaned or no parent)
 
         # Kahn's algorithm: start with nodes with in-degree 0
-        # Results are sorted by metadata.creation_time
+        # Results are sorted by metadata.creation_time, those without creation time are last
         queue = [
-            (metadata.creation_time, metadata)
+            (metadata.creation_time or datetime.datetime.max, metadata)
             for metadata in run_metadata
             if in_degree[SourceRunId(metadata.run_id)] == 0
         ]
@@ -141,7 +142,11 @@ def _topological_sort_runs(
             if source_run_id in parent_to_children:
                 for child_metadata in parent_to_children[source_run_id]:
                     heapq.heappush(
-                        queue, (child_metadata.creation_time, child_metadata)
+                        queue,
+                        (
+                            child_metadata.creation_time or datetime.datetime.max,
+                            child_metadata,
+                        ),
                     )
 
         # Check for cycles (shouldn't happen in Neptune, but defensive)
@@ -157,6 +162,9 @@ def _topological_sort_runs(
                 for metadata in run_metadata
                 if SourceRunId(metadata.run_id) not in result_ids
             ]
+            remaining_metadata.sort(
+                key=lambda x: (x.creation_time or datetime.datetime.max, x.run_id)
+            )
             result.extend(remaining_metadata)
 
         return result

tests/unit/test_loader_manager.py

@@ -204,8 +204,10 @@ def test_topological_sort_creation_time_order():
     # Setup: two runs with different creation times
     run1_id = "RUN-1"
     run2_id = "RUN-2"
+    run3_id = "RUN-3"
     creation_time1 = datetime.datetime(2025, 1, 1, 12, 0, 0)
-    creation_time2 = datetime.datetime(2025, 1, 1, 12, 0, 1)
+    creation_time2 = None
+    creation_time3 = datetime.datetime(2025, 1, 1, 12, 0, 2)
 
     run1_metadata = RunMetadata(
         project_id="project1",
@@ -227,9 +229,22 @@ def test_topological_sort_creation_time_order():
         creation_time=creation_time2,
     )
 
+    run3_metadata = RunMetadata(
+        project_id="project1",
+        run_id=run3_id,
+        custom_run_id=None,
+        experiment_name=None,
+        parent_source_run_id=None,
+        fork_step=None,
+        creation_time=creation_time3,
+    )
     # Mock reader
-    mock_reader.list_run_files.return_value = [run2_id, run1_id]
-    mock_reader.read_run_metadata.side_effect = [run2_metadata, run1_metadata]
+    mock_reader.list_run_files.return_value = [run3_id, run2_id, run1_id]
+    mock_reader.read_run_metadata.side_effect = [
+        run3_metadata,
+        run2_metadata,
+        run1_metadata,
+    ]
     mock_reader.read_run_data.return_value = iter([pa.table({"col": [1]})])
 
     # Track creation order
@@ -248,9 +263,10 @@ def track_create_run(*args, **kwargs):
     manager._load_project(project_dir, runs=None)
 
     # Verify: runs should be processed in creation time order
-    assert len(created_runs) == 2
+    assert len(created_runs) == 3
     assert created_runs[0] == run1_id
-    assert created_runs[1] == run2_id
+    assert created_runs[1] == run3_id
+    assert created_runs[2] == run2_id
 
 
 def test_topological_sort_multiple_children():