orchestrata_ultimate.kernel.pipeline_checkpoint

Reference › orchestrata_ultimate › kernel › pipeline_checkpoint

The checkpoint store that lets an interrupted pipeline run resume without re-prompting the LLM stages that already finished. PipelineCheckpoint saves each completed stage's output, keyed by thread id and stage alias, and StageCheckpoint is the stored record (thread, stage, subject, output, timestamp, optional work type). It serves SPEC-003 requirement FR-024 (REQ-003-FR24-002, -003). pipeline owns one, pipeline_session fills it, and the orchestrator's pipeline_escalation and pipeline_runner modules and the session pipeline CLI import it.

save_stage requires a thread id and a stage. It keeps the full output in an in-process dictionary and writes a JSON record to a Redis hash named from the thread id; when the encoded record exceeds an inline size limit, _spill_to_disk writes it to a file and stores only a pointer. The kernel passes the postman's Redis client; without one, the store uses the in-process mirror and disk. The spill directory comes from the checkpoint_dir argument or PIPELINE_CHECKPOINT_DIR, defaulting to a pipeline_checkpoints folder. get_stage reads the in-process mirror first, then Redis, resolving disk pointers. completed_stages and resume_state list the stages and return a mapping of stage to output; stored_work_type, has_stage and clear complete the interface. A Redis failure on save is logged and not raised, so the in-process copy still serves the live run. PipelineRunner.resume uses resume_state to find the last completed stage and re-dispatch from the next one.

At a glance

Class diagram

Dependency diagram

API

Classes

StageCheckpoint (dataclass)

Decorators: dataclass, traces(spec='003', story='FR-024', reqs=['REQ-003-FR24-003']) · Line 37

A single persisted stage output.

PipelineCheckpoint

Decorators: traces(spec='003', story='FR-024', reqs=['REQ-003-FR24-002', 'REQ-003-FR24-003']) · Line 57

Persists completed stage outputs and resumes from them.

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/kernel/pipeline_checkpoint.py

"""Pipeline checkpoint / resume — reuse completed-stage outputs.

LLM stages (BA/PM/Ar/Cd/QA) are not replayable: a resumed run must reuse the
output a stage already produced rather than re-invoking the model. This module
persists each completed stage's output keyed by ``thread_id`` so a resume can
skip completed stages.

Storage strategy (from the plan):
  - small state  → Redis hash ``pipeline:checkpoint:<thread_id>`` (field=stage)
  - large blobs  → disk under ``pipeline_checkpoints/<thread_id>/<stage>.json``
                   with a small pointer record kept in Redis
When Redis is unavailable the checkpoint degrades to an in-process dict (still
correct within a single runner lifetime) plus the on-disk blobs.
"""

import json
import logging
import os
import threading
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from orchestrata_ultimate.core.base.traces import traces

_logger = logging.getLogger(__name__)

# Redis key prefix for per-thread checkpoint hashes.
_KEY_PREFIX = "pipeline:checkpoint:"
# Outputs whose JSON exceeds this many bytes are spilled to disk with a pointer.
_INLINE_MAX_BYTES = 16 * 1024


@dataclass
@traces(spec="003", story="FR-024", reqs=["REQ-003-FR24-003"])
class StageCheckpoint:
    """A single persisted stage output.

    Attributes:
        thread_id: Pipeline thread the stage belongs to.
        stage: Volk alias / stage that produced the output.
        subject: CXP subject of the transition that produced it.
        output: The stage output payload (must be JSON-serialisable).
        timestamp: When the checkpoint was written (UTC, ISO-8601).
    """

    thread_id: str
    stage: str
    subject: str
    output: dict[str, Any]
    timestamp: str = field(default_factory=lambda: datetime.now(UTC).isoformat())
    work_type: str | None = None


@traces(spec="003", story="FR-024", reqs=["REQ-003-FR24-002", "REQ-003-FR24-003"])
class PipelineCheckpoint:
    """Persists completed stage outputs and resumes from them.

    Small outputs live inline in a Redis hash; large outputs spill to disk with
    a Redis pointer. Thread-safe for the free-threaded kernel.
    """

    def __init__(
        self,
        redis_client: Any = None,
        checkpoint_dir: Path | str | None = None,
    ) -> None:
        """Initialise the checkpoint store.

        Args:
            redis_client: Sync Redis client (the kernel passes the postman's).
                When None, falls back to an in-process dict + disk.
            checkpoint_dir: Root for spilled blobs. Defaults to
                ``$PIPELINE_CHECKPOINT_DIR`` or ``./pipeline_checkpoints``.
        """
        self._redis = redis_client
        self._lock = threading.RLock()
        # In-process mirror — authoritative when Redis is absent, and a fast
        # read path / fallback when it is present.
        self._memory: dict[str, dict[str, StageCheckpoint]] = {}
        if checkpoint_dir is None:
            checkpoint_dir = os.environ.get(
                "PIPELINE_CHECKPOINT_DIR", "pipeline_checkpoints"
            )
        self._dir = Path(checkpoint_dir)

    def _key(self, thread_id: str) -> str:
        return f"{_KEY_PREFIX}{thread_id}"

    def _blob_path(self, thread_id: str, stage: str) -> Path:
        safe_stage = stage.replace(":", "-").replace("/", "-") or "stage"
        return self._dir / thread_id / f"{safe_stage}.json"

    @traces(spec="003", story="FR-024", reqs=["REQ-003-FR24-003"])
    def save_stage(
        self,
        *,
        thread_id: str,
        stage: str,
        subject: str,
        output: dict[str, Any],
        work_type: str | None = None,
    ) -> StageCheckpoint:
        """Persist a completed stage's output.

        Small outputs are stored inline in the Redis hash; large outputs are
        written to disk and a pointer record (``{"_blob": <path>}``) is stored
        instead. The in-process mirror always keeps the full output.

        Args:
            thread_id: Pipeline thread id.
            stage: Producing stage / volk alias.
            subject: Transition subject.
            output: JSON-serialisable stage output.

        Returns:
            The :class:`StageCheckpoint` that was stored.
        """
        if not thread_id or not stage:
            raise ValueError("thread_id and stage are required to checkpoint")

        cp = StageCheckpoint(
            thread_id=thread_id,
            stage=stage,
            subject=subject,
            output=dict(output),
            work_type=work_type,
        )
        with self._lock:
            self._memory.setdefault(thread_id, {})[stage] = cp

        record = {
            "thread_id": thread_id,
            "stage": stage,
            "subject": subject,
            "output": cp.output,
            "timestamp": cp.timestamp,
            "work_type": cp.work_type,
        }
        encoded = json.dumps(record, default=str)

        if len(encoded.encode("utf-8")) > _INLINE_MAX_BYTES:
            encoded = self._spill_to_disk(thread_id, stage, record)

        if self._redis is not None:
            try:
                self._redis.hset(self._key(thread_id), stage, encoded)
            except Exception as exc:
                _logger.warning("Checkpoint Redis save failed: %s", exc)
        return cp

    def _spill_to_disk(self, thread_id: str, stage: str, record: dict[str, Any]) -> str:
        """Write a large stage record to disk; return a pointer JSON string."""
        path = self._blob_path(thread_id, stage)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(record, default=str), encoding="utf-8")
        return json.dumps({"_blob": str(path)})

    @traces(spec="003", story="FR-024", reqs=["REQ-003-FR24-003"])
    def get_stage(self, thread_id: str, stage: str) -> dict[str, Any] | None:
        """Return a completed stage's output, or None if not checkpointed.

        Reads the in-process mirror first, then Redis (resolving disk pointers).

        Args:
            thread_id: Pipeline thread id.
            stage: Stage / volk alias to fetch.
        """
        with self._lock:
            mem = self._memory.get(thread_id, {}).get(stage)
        if mem is not None:
            return mem.output

        record = self._load_record(thread_id, stage)
        return record.get("output") if record else None

    def _load_record(self, thread_id: str, stage: str) -> dict[str, Any] | None:
        """Load a stage record from Redis, resolving a disk pointer if present."""
        if self._redis is None:
            return None
        try:
            raw = self._redis.hget(self._key(thread_id), stage)
        except Exception as exc:
            _logger.warning("Checkpoint Redis read failed: %s", exc)
            return None
        if not raw:
            return None
        if isinstance(raw, bytes):
            raw = raw.decode("utf-8")
        try:
            data: dict[str, Any] = json.loads(raw)
        except json.JSONDecodeError:
            return None
        if isinstance(data, dict) and "_blob" in data:
            return self._load_blob(data["_blob"])
        return data

    def _load_blob(self, blob_path: str) -> dict[str, Any] | None:
        """Load a spilled stage record from disk."""
        try:
            blob: dict[str, Any] = json.loads(
                Path(blob_path).read_text(encoding="utf-8")
            )
            return blob
        except (OSError, json.JSONDecodeError) as exc:
            _logger.warning("Checkpoint blob read failed (%s): %s", blob_path, exc)
            return None

    @traces(spec="003", story="FR-024", reqs=["REQ-003-FR24-002"])
    def completed_stages(self, thread_id: str) -> list[str]:
        """Return the stages already checkpointed for *thread_id*.

        Args:
            thread_id: Pipeline thread id.

        Returns:
            Stage aliases that have a stored output (sorted for stability).
        """
        with self._lock:
            stages = set(self._memory.get(thread_id, {}).keys())
        if self._redis is not None:
            try:
                keys = self._redis.hkeys(self._key(thread_id))
                for k in keys:
                    stages.add(k.decode("utf-8") if isinstance(k, bytes) else k)
            except Exception as exc:
                _logger.warning("Checkpoint Redis hkeys failed: %s", exc)
        return sorted(stages)

    @traces(spec="003", story="FR-024", reqs=["REQ-003-FR24-003"])
    def resume_state(self, thread_id: str) -> dict[str, dict[str, Any]]:
        """Return every completed stage's output for a resume.

        The runner uses this on resume to skip stages that already produced
        output (LLM stages are not replayable) and reuse their result.

        Args:
            thread_id: Pipeline thread id.

        Returns:
            ``{stage: output}`` for every checkpointed stage.
        """
        state: dict[str, dict[str, Any]] = {}
        for stage in self.completed_stages(thread_id):
            output = self.get_stage(thread_id, stage)
            if output is not None:
                state[stage] = output
        return state

    @traces(spec="003", story="FR-024", reqs=["REQ-003-FR24-003"])
    def stored_work_type(self, thread_id: str) -> str | None:
        """Return the original run's work_type from any checkpointed stage.

        Older checkpoint records (persisted before work_type tracking) lack
        the field entirely, so this returns ``None`` when no checkpointed
        stage for *thread_id* carries one — callers fall back to their own
        default in that case.

        Args:
            thread_id: Pipeline thread id.

        Returns:
            The persisted work_type, or None if never recorded.
        """
        for stage in self.completed_stages(thread_id):
            with self._lock:
                mem = self._memory.get(thread_id, {}).get(stage)
            if mem is not None and mem.work_type is not None:
                return mem.work_type
            record = self._load_record(thread_id, stage)
            if record and record.get("work_type") is not None:
                wt = record["work_type"]
                return wt if isinstance(wt, str) else None
        return None

    @traces(spec="003", story="FR-024", reqs=["REQ-003-FR24-003"])
    def has_stage(self, thread_id: str, stage: str) -> bool:
        """Return True if *stage* has a checkpointed output for *thread_id*."""
        return self.get_stage(thread_id, stage) is not None

    @traces(spec="000", reqs=["REQ-000-053"])
    def clear(self, thread_id: str) -> None:
        """Drop every checkpoint for a thread (Redis hash, disk, memory).

        Args:
            thread_id: Pipeline thread id to clear.
        """
        with self._lock:
            self._memory.pop(thread_id, None)
        if self._redis is not None:
            try:
                self._redis.delete(self._key(thread_id))
            except Exception as exc:
                _logger.warning("Checkpoint Redis clear failed: %s", exc)
        blob_dir = self._dir / thread_id
        if blob_dir.is_dir():
            for child in blob_dir.glob("*.json"):
                try:
                    child.unlink()
                except OSError:
                    pass