Improve runtime profiling storage resilience

Author: tobiasraabe

src/_pytask/journal.py

@@ -28,19 +28,25 @@ def append(self, payload: msgspec.Struct) -> None:
             journal_file.write(msgspec.json.encode(payload) + b"\n")
 
     def read(self) -> list[T]:
-        """Read entries, clearing the journal on decode errors."""
+        """Read entries, keeping valid entries on decode errors."""
         if not self.path.exists():
             return []
 
         entries: list[T] = []
-        for line in self.path.read_bytes().splitlines():
-            if not line.strip():
+        data = self.path.read_bytes()
+        offset = 0
+        for line in data.splitlines(keepends=True):
+            stripped = line.strip()
+            if not stripped:
+                offset += len(line)
                 continue
             try:
-                entries.append(msgspec.json.decode(line, type=self.type_))
+                entries.append(msgspec.json.decode(stripped, type=self.type_))
             except msgspec.DecodeError:
-                self.delete()
-                return []
+                with self.path.open("rb+") as journal_file:
+                    journal_file.truncate(offset)
+                return entries
+            offset += len(line)
         return entries
 
     def delete(self) -> None:

src/_pytask/profile.py

@@ -99,6 +99,7 @@ def pytask_unconfigure(self, session: Session) -> None:
         self.runtime_state.flush()
 
 
+@dataclass
 class DurationNameSpace:
     """A namespace for adding durations to the profile."""
 

src/_pytask/runtime_store.py

@@ -15,8 +15,6 @@
 
     from _pytask.node_protocols import PTask
 
-CURRENT_RUNTIME_VERSION = "1"
-
 
 class _RuntimeEntry(msgspec.Struct):
     id: str
@@ -25,12 +23,10 @@ class _RuntimeEntry(msgspec.Struct):
 
 
 class _RuntimeFile(msgspec.Struct, forbid_unknown_fields=False):
-    runtime_version: str = msgspec.field(name="runtime-version")
     task: list[_RuntimeEntry] = msgspec.field(default_factory=list)
 
 
-class _RuntimeJournalEntry(msgspec.Struct):
-    runtime_version: str = msgspec.field(name="runtime-version")
+class _RuntimeJournalEntry(msgspec.Struct, forbid_unknown_fields=False):
     id: str
     date: float
     duration: float
@@ -56,10 +52,6 @@ def _read_runtimes(path: Path) -> _RuntimeFile | None:
     except msgspec.DecodeError:
         path.unlink()
         return None
-
-    if data.runtime_version != CURRENT_RUNTIME_VERSION:
-        path.unlink()
-        return None
     return data
 
 
@@ -73,12 +65,7 @@ def _write_runtimes(path: Path, runtimes: _RuntimeFile) -> None:
 def _read_journal(
     journal: JsonlJournal[_RuntimeJournalEntry],
 ) -> list[_RuntimeJournalEntry]:
-    entries = journal.read()
-    for entry in entries:
-        if entry.runtime_version != CURRENT_RUNTIME_VERSION:
-            journal.delete()
-            return []
-    return entries
+    return journal.read()
 
 
 def _apply_journal(
@@ -92,7 +79,6 @@ def _apply_journal(
             id=entry.id, date=entry.date, duration=entry.duration
         )
     return _RuntimeFile(
-        runtime_version=CURRENT_RUNTIME_VERSION,
         task=list(index.values()),
     )
 
@@ -116,7 +102,6 @@ def from_root(cls, root: Path) -> RuntimeState:
         journal_entries = _read_journal(journal)
         if existing is None:
             runtimes = _RuntimeFile(
-                runtime_version=CURRENT_RUNTIME_VERSION,
                 task=[],
             )
             runtimes = _apply_journal(runtimes, journal_entries)
@@ -137,12 +122,10 @@ def update_task(self, task: PTask, start: float, end: float) -> None:
         entry = _RuntimeEntry(id=task_id, date=start, duration=end - start)
         self._index[entry.id] = entry
         self.runtimes = _RuntimeFile(
-            runtime_version=CURRENT_RUNTIME_VERSION,
             task=list(self._index.values()),
         )
         self._rebuild_index()
         journal_entry = _RuntimeJournalEntry(
-            runtime_version=CURRENT_RUNTIME_VERSION,
             id=entry.id,
             date=entry.date,
             duration=entry.duration,

tests/test_runtime_store.py

@@ -58,3 +58,22 @@ def test_runtime_state_flushes_journal(tmp_path):
 
     reloaded = RuntimeState.from_root(tmp_path)
     assert reloaded.get_duration(task) == pytest.approx(3.5)
+
+
+def test_runtime_state_recovers_from_corrupt_journal(tmp_path):
+    tmp_path.joinpath(".pytask").mkdir()
+    task_a = DummyTask(name="task_a")
+    task_b = DummyTask(name="task_b")
+
+    state = RuntimeState.from_root(tmp_path)
+    state.update_task(task_a, 1.0, 3.0)
+    state.update_task(task_b, 2.0, 6.0)
+
+    journal_path = tmp_path / ".pytask" / "runtimes.journal"
+    with journal_path.open("ab") as journal_file:
+        journal_file.write(b'{"id": "corrupt"')
+
+    recovered = RuntimeState.from_root(tmp_path)
+    assert recovered.get_duration(task_a) == pytest.approx(2.0)
+    assert recovered.get_duration(task_b) == pytest.approx(4.0)
+    assert b'"corrupt"' not in journal_path.read_bytes()