Move runtimes to json journal store

Author: tobiasraabe

CHANGELOG.md

@@ -7,7 +7,9 @@ releases are available on [PyPI](https://pypi.org/project/pytask) and
 
 ## Unreleased
 
-- Nothing yet.
+- Move runtime profiling persistence from SQLite to a JSON snapshot plus append-only
+  journal in `.pytask/`, keeping runtime data resilient to crashes and compacted on
+  normal build exits.
 
 ## 0.5.8 - 2025-12-30
 

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
     "attrs>=21.3.0",
     "click>=8.1.8,!=8.2.0",
     "click-default-group>=1.2.4",
+    "msgspec>=0.18.6",
     "networkx>=2.4.0",
     "optree>=0.9.0",
     "packaging>=23.0.0",

src/_pytask/journal.py

@@ -0,0 +1,41 @@
+"""Helpers for append-only JSONL journals."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from typing import TypeVar
+
+import msgspec
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+T = TypeVar("T")
+
+
+def append_jsonl(path: Path, payload: msgspec.Struct) -> None:
+    """Append a JSON line to the journal."""
+    with path.open("ab") as journal_file:
+        journal_file.write(msgspec.json.encode(payload) + b"\n")
+
+
+def read_jsonl(path: Path, *, type_: type[T]) -> list[T]:
+    """Read JSONL entries from a journal, stopping at the first invalid line."""
+    if not path.exists():
+        return []
+
+    entries: list[T] = []
+    for line in path.read_bytes().splitlines():
+        if not line.strip():
+            continue
+        try:
+            entries.append(msgspec.json.decode(line, type=type_))
+        except msgspec.DecodeError:
+            break
+    return entries
+
+
+def delete_if_exists(path: Path) -> None:
+    """Delete a file if it exists."""
+    if path.exists():
+        path.unlink()

src/_pytask/profile.py

@@ -13,16 +13,12 @@
 
 import click
 from rich.table import Table
-from sqlalchemy.orm import Mapped
-from sqlalchemy.orm import mapped_column
 
 from _pytask.click import ColoredCommand
 from _pytask.click import EnumChoice
 from _pytask.console import console
 from _pytask.console import format_task_name
 from _pytask.dag import create_dag
-from _pytask.database_utils import BaseTable
-from _pytask.database_utils import DatabaseSession
 from _pytask.exceptions import CollectionError
 from _pytask.exceptions import ConfigurationError
 from _pytask.node_protocols import PPathNode
@@ -31,6 +27,7 @@
 from _pytask.outcomes import TaskOutcome
 from _pytask.pluginmanager import hookimpl
 from _pytask.pluginmanager import storage
+from _pytask.runtime_store import RuntimeState
 from _pytask.session import Session
 from _pytask.traceback import Traceback
 
@@ -48,16 +45,6 @@ class _ExportFormats(enum.Enum):
     CSV = "csv"
 
 
-class Runtime(BaseTable):
-    """Record of runtimes of tasks."""
-
-    __tablename__ = "runtime"
-
-    task: Mapped[str] = mapped_column(primary_key=True)
-    date: Mapped[float]
-    duration: Mapped[float]
-
-
 @hookimpl(tryfirst=True)
 def pytask_extend_command_line_interface(cli: click.Group) -> None:
     """Extend the command line interface."""
@@ -67,6 +54,7 @@ def pytask_extend_command_line_interface(cli: click.Group) -> None:
 @hookimpl
 def pytask_post_parse(config: dict[str, Any]) -> None:
     """Register the export option."""
+    config["runtime_state"] = RuntimeState.from_root(config["root"])
     config["pm"].register(ExportNameSpace)
     config["pm"].register(DurationNameSpace)
     config["pm"].register(FileSizeNameSpace)
@@ -83,26 +71,16 @@ def pytask_execute_task(task: PTask) -> Generator[None, None, None]:
 
 
 @hookimpl
-def pytask_execute_task_process_report(report: ExecutionReport) -> None:
-    """Store runtime of successfully finishing tasks in database."""
+def pytask_execute_task_process_report(
+    session: Session, report: ExecutionReport
+) -> None:
+    """Store runtime of successfully finishing tasks."""
     task = report.task
     duration = task.attributes.get("duration")
     if report.outcome == TaskOutcome.SUCCESS and duration is not None:
-        _create_or_update_runtime(task.signature, *duration)
-
-
-def _create_or_update_runtime(task_signature: str, start: float, end: float) -> None:
-    """Create or update a runtime entry."""
-    with DatabaseSession() as session:
-        runtime = session.get(Runtime, task_signature)
-
-        if not runtime:
-            session.add(Runtime(task=task_signature, date=start, duration=end - start))
-        else:
-            for attr, val in (("date", start), ("duration", end - start)):
-                setattr(runtime, attr, val)
-
-        session.commit()
+        runtime_state = session.config.get("runtime_state")
+        if runtime_state is not None:
+            runtime_state.update_task(task, *duration)
 
 
 @click.command(cls=ColoredCommand)
@@ -189,20 +167,23 @@ class DurationNameSpace:
     @staticmethod
     @hookimpl
     def pytask_profile_add_info_on_task(
-        tasks: list[PTask], profile: dict[str, dict[str, Any]]
+        session: Session, tasks: list[PTask], profile: dict[str, dict[str, Any]]
     ) -> None:
         """Add the runtime for tasks to the profile."""
-        runtimes = _collect_runtimes(tasks)
+        runtimes = _collect_runtimes(session, tasks)
         for name, duration in runtimes.items():
             profile[name]["Duration (in s)"] = round(duration, 2)
 
 
-def _collect_runtimes(tasks: list[PTask]) -> dict[str, float]:
+def _collect_runtimes(session: Session, tasks: list[PTask]) -> dict[str, float]:
     """Collect runtimes."""
-    with DatabaseSession() as session:
-        runtimes = [session.get(Runtime, task.signature) for task in tasks]
+    runtime_state = session.config.get("runtime_state")
+    if runtime_state is None:
+        return {}
     return {
-        task.name: r.duration for task, r in zip(tasks, runtimes, strict=False) if r
+        task.name: duration
+        for task in tasks
+        if (duration := runtime_state.get_duration(task)) is not None
     }
 
 
@@ -313,3 +294,16 @@ def _get_info_names(profile: dict[str, dict[str, Any]]) -> list[str]:
     base: set[str] = set()
     info_names: list[str] = sorted(base.union(*(set(val) for val in profile.values())))
     return info_names
+
+
+@hookimpl
+def pytask_unconfigure(session: Session) -> None:
+    """Flush runtime information on normal build exits."""
+    if session.config.get("command") != "build":
+        return
+    if session.config.get("dry_run") or session.config.get("explain"):
+        return
+    runtime_state = session.config.get("runtime_state")
+    if runtime_state is None:
+        return
+    runtime_state.flush()

src/_pytask/runtime_store.py

@@ -0,0 +1,179 @@
+"""Runtime storage with an append-only journal."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from dataclasses import field
+from typing import TYPE_CHECKING
+
+import msgspec
+from packaging.version import Version
+
+from _pytask.journal import append_jsonl
+from _pytask.journal import delete_if_exists
+from _pytask.journal import read_jsonl
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from _pytask.node_protocols import PTask
+
+CURRENT_RUNTIME_VERSION = "1"
+
+
+class RuntimeStoreError(Exception):
+    """Raised when reading or writing runtime files fails."""
+
+
+class RuntimeStoreVersionError(RuntimeStoreError):
+    """Raised when a runtime file version is not supported."""
+
+
+class _RuntimeEntry(msgspec.Struct):
+    id: str
+    date: float
+    duration: float
+
+
+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")
+    id: str
+    date: float
+    duration: float
+
+
+def _runtimes_path(root: Path) -> Path:
+    return root / ".pytask" / "runtimes.json"
+
+
+def _journal_path(path: Path) -> Path:
+    return path.with_suffix(".journal")
+
+
+def _read_runtimes(path: Path) -> _RuntimeFile | None:
+    if not path.exists():
+        return None
+    try:
+        data = msgspec.json.decode(path.read_bytes(), type=_RuntimeFile)
+    except msgspec.DecodeError:
+        msg = "Runtime file has invalid format."
+        raise RuntimeStoreError(msg) from None
+
+    if Version(data.runtime_version) != Version(CURRENT_RUNTIME_VERSION):
+        msg = (
+            f"Unsupported runtime-version {data.runtime_version!r}. "
+            f"Current version is {CURRENT_RUNTIME_VERSION}."
+        )
+        raise RuntimeStoreVersionError(msg)
+    return data
+
+
+def _write_runtimes(path: Path, runtimes: _RuntimeFile) -> None:
+    data = msgspec.json.encode(runtimes)
+    tmp = path.with_suffix(f"{path.suffix}.tmp")
+    tmp.write_bytes(data)
+    tmp.replace(path)
+
+
+def _read_journal(path: Path) -> list[_RuntimeJournalEntry]:
+    journal_path = _journal_path(path)
+    entries = read_jsonl(journal_path, type_=_RuntimeJournalEntry)
+    for entry in entries:
+        if Version(entry.runtime_version) != Version(CURRENT_RUNTIME_VERSION):
+            msg = (
+                f"Unsupported runtime-version {entry.runtime_version!r}. "
+                f"Current version is {CURRENT_RUNTIME_VERSION}."
+            )
+            raise RuntimeStoreVersionError(msg)
+    return entries
+
+
+def _apply_journal(
+    runtimes: _RuntimeFile, entries: list[_RuntimeJournalEntry]
+) -> _RuntimeFile:
+    if not entries:
+        return runtimes
+    index = {entry.id: entry for entry in runtimes.task}
+    for entry in entries:
+        index[entry.id] = _RuntimeEntry(
+            id=entry.id, date=entry.date, duration=entry.duration
+        )
+    return _RuntimeFile(
+        runtime_version=CURRENT_RUNTIME_VERSION,
+        task=list(index.values()),
+    )
+
+
+def _build_task_id(task: PTask) -> str:
+    return task.name
+
+
+@dataclass
+class RuntimeState:
+    path: Path
+    runtimes: _RuntimeFile
+    _index: dict[str, _RuntimeEntry] = field(init=False, default_factory=dict)
+    _dirty: bool = field(init=False, default=False)
+
+    def __post_init__(self) -> None:
+        self._rebuild_index()
+
+    @classmethod
+    def from_root(cls, root: Path) -> RuntimeState:
+        path = _runtimes_path(root)
+        existing = _read_runtimes(path)
+        journal_entries = _read_journal(path)
+        if existing is None:
+            runtimes = _RuntimeFile(
+                runtime_version=CURRENT_RUNTIME_VERSION,
+                task=[],
+            )
+            runtimes = _apply_journal(runtimes, journal_entries)
+            state = cls(path=path, runtimes=runtimes)
+        else:
+            runtimes = _apply_journal(existing, journal_entries)
+            state = cls(path=path, runtimes=runtimes)
+
+        if journal_entries:
+            state._dirty = True
+        return state
+
+    def _rebuild_index(self) -> None:
+        self._index = {entry.id: entry for entry in self.runtimes.task}
+
+    def update_task(self, task: PTask, start: float, end: float) -> None:
+        task_id = _build_task_id(task)
+        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,
+        )
+        append_jsonl(_journal_path(self.path), journal_entry)
+        self._dirty = True
+
+    def get_duration(self, task: PTask) -> float | None:
+        task_id = _build_task_id(task)
+        entry = self._index.get(task_id)
+        if entry is None:
+            return None
+        return entry.duration
+
+    def flush(self) -> None:
+        if not self._dirty:
+            return
+        _write_runtimes(self.path, self.runtimes)
+        delete_if_exists(_journal_path(self.path))
+        self._dirty = False

src/pytask/__init__.py

@@ -66,7 +66,6 @@
 from _pytask.pluginmanager import get_plugin_manager
 from _pytask.pluginmanager import hookimpl
 from _pytask.pluginmanager import storage
-from _pytask.profile import Runtime
 from _pytask.reports import CollectionReport
 from _pytask.reports import DagReport
 from _pytask.reports import ExecutionReport
@@ -123,7 +122,6 @@
     "PytaskError",
     "PythonNode",
     "ResolvingDependenciesError",
-    "Runtime",
     "Session",
     "ShowCapture",
     "Skipped",