deer-flow/backend/scripts/benchmark/checkpoint/bench_channels.py

#!/usr/bin/env python3
"""Benchmark DeerFlow's full and DeltaChannel checkpoint message storage.

The public CLI is a controller. Every benchmark case runs in a fresh child
process and, for SQLite, a fresh database. This mirrors the restart-required
checkpoint mode boundary and prevents one mode's graph/channel caches from
warming the other.

Examples::

    PYTHONPATH=. uv run python scripts/benchmark/checkpoint/bench_channels.py \
        --updates 10,100,500,999,1000,1001,2000 \
        --payload-bytes 128,4096 \
        --output checkpoint-bench.jsonl

    PYTHONPATH=. uv run python scripts/benchmark/checkpoint/bench_channels.py \
        --backends sqlite --updates 1000 --payload-bytes 128 \
        --repetitions 7 --output snapshot-boundary.jsonl

The controller suppresses matrix pairs whose estimated cumulative full-mode
message payload exceeds ``--max-estimated-full-bytes``. Both modes are skipped
as a pair so every emitted full/delta result remains comparable. Use
``--allow-large-cases`` only on a machine provisioned for the resulting disk
and memory use.
"""

from __future__ import annotations

import argparse
import gc
import importlib.metadata
import json
import os
import platform
import sys
import tempfile
import time
from collections.abc import Callable
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Annotated, Any, Literal, TypedDict

from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, HumanMessage
from langgraph.channels import DeltaChannel
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages

from deerflow.agents.thread_state import merge_message_writes
from deerflow.runtime.checkpoint_mode import inject_checkpoint_mode
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor

sys.path.insert(0, str(Path(__file__).resolve().parent))
import checkpoint_bench_common as _common  # noqa: E402

_GIT_SHA_ENV = _common.GIT_SHA_ENV
_parse_positive_int_csv = _common.parse_positive_int_csv
_parse_choice_csv = _common.parse_choice_csv
_canonical_messages_digest = _common.canonical_messages_digest
_percentile = _common.percentile
_window_median = _common.window_median
_resolve_git_sha = _common.resolve_git_sha
_safe_error = _common.safe_error
_file_size = _common.file_size
_peak_rss_bytes = _common.peak_rss_bytes

Mode = Literal["full", "delta"]
Backend = Literal["memory", "sqlite"]

SCHEMA_VERSION = 1
BENCHMARK_VERSION = 1
PRODUCTION_SNAPSHOT_FREQUENCY = 1000
DEFAULT_MAX_ESTIMATED_FULL_BYTES = 1024**3
_MODES: tuple[Mode, ...] = ("full", "delta")
_BACKENDS: tuple[Backend, ...] = ("memory", "sqlite")
_STORAGE_STAT_FIELDS = (
    "logical_checkpoint_bytes",
    "logical_write_bytes",
    "checkpoint_rows",
    "write_rows",
)


class _FullBenchmarkState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]


class _DeltaBenchmarkState(TypedDict):
    messages: Annotated[
        list[AnyMessage],
        DeltaChannel(merge_message_writes, snapshot_frequency=PRODUCTION_SNAPSHOT_FREQUENCY),
    ]


@dataclass(frozen=True)
class BenchmarkCase:
    mode: Mode
    backend: Backend
    update_count: int
    payload_bytes: int
    repetition: int
    seed: int
    scenario: str = "append"
    snapshot_frequency: int = PRODUCTION_SNAPSHOT_FREQUENCY

    def __post_init__(self) -> None:
        if self.mode not in _MODES:
            raise ValueError(f"unsupported mode: {self.mode!r}")
        if self.backend not in _BACKENDS:
            raise ValueError(f"unsupported backend: {self.backend!r}")
        if self.update_count <= 0:
            raise ValueError("update_count must be positive")
        if self.payload_bytes <= 0:
            raise ValueError("payload_bytes must be positive")
        if self.repetition < 0:
            raise ValueError("repetition must be non-negative")
        if self.snapshot_frequency != PRODUCTION_SNAPSHOT_FREQUENCY:
            raise ValueError(f"Phase 1 supports only production snapshot_frequency={PRODUCTION_SNAPSHOT_FREQUENCY}")
        if self.scenario != "append":
            raise ValueError(f"unsupported scenario: {self.scenario!r}")


def _expand_cases(
    *,
    modes: list[str],
    backends: list[str],
    update_counts: list[int],
    payload_bytes: list[int],
    repetitions: int,
    seed: int,
) -> list[BenchmarkCase]:
    """Build a matrix with alternating mode order to reduce order bias."""
    cases: list[BenchmarkCase] = []
    for repetition in range(repetitions):
        for backend in backends:
            for payload in payload_bytes:
                for update_index, update_count in enumerate(update_counts):
                    ordered_modes = list(modes)
                    if (repetition + update_index) % 2 == 1:
                        ordered_modes.reverse()
                    for mode in ordered_modes:
                        cases.append(
                            BenchmarkCase(
                                mode=mode,  # type: ignore[arg-type]
                                backend=backend,  # type: ignore[arg-type]
                                update_count=update_count,
                                payload_bytes=payload,
                                repetition=repetition,
                                seed=seed,
                            )
                        )
    return cases


def _estimated_full_payload_bytes(case: BenchmarkCase) -> int:
    """Lower-bound cumulative message content serialized by full mode."""
    return case.payload_bytes * case.update_count * (case.update_count + 1) // 2


def _filter_oversized_pairs(cases: list[BenchmarkCase], *, max_bytes: int | None) -> tuple[list[BenchmarkCase], list[BenchmarkCase]]:
    if max_bytes is None:
        return cases, []
    group_fields = ("backend", "scenario", "snapshot_frequency", "update_count", "payload_bytes", "repetition")
    oversized_keys = {tuple(getattr(case, field) for field in group_fields) for case in cases if case.mode == "full" and _estimated_full_payload_bytes(case) > max_bytes}
    kept = [case for case in cases if tuple(getattr(case, field) for field in group_fields) not in oversized_keys]
    skipped = [case for case in cases if tuple(getattr(case, field) for field in group_fields) in oversized_keys]
    return kept, skipped


def _message_for_update(index: int, payload_bytes: int) -> BaseMessage:
    content = "x" * payload_bytes
    message_id = f"bench-message-{index:08d}"
    if index % 2 == 0:
        return HumanMessage(id=message_id, content=content)
    return AIMessage(id=message_id, content=content)


def _noop(_state: dict[str, Any]) -> dict[str, Any]:
    return {}


def _build_graph(mode: Mode, saver: Any) -> Any:
    schema = _DeltaBenchmarkState if mode == "delta" else _FullBenchmarkState
    builder = StateGraph(schema)
    builder.add_node("noop", _noop)
    builder.set_entry_point("noop")
    builder.set_finish_point("noop")
    return builder.compile(checkpointer=saver)


def _config(case: BenchmarkCase) -> dict[str, Any]:
    config: dict[str, Any] = {
        "configurable": {
            "thread_id": f"checkpoint-bench-{case.seed}-{case.repetition}",
        }
    }
    inject_checkpoint_mode(config, case.mode)
    return config


def _base_row(case: BenchmarkCase) -> dict[str, Any]:
    git_sha = os.environ.get(_GIT_SHA_ENV) or _resolve_git_sha()
    try:
        langgraph_version = importlib.metadata.version("langgraph")
    except importlib.metadata.PackageNotFoundError:
        langgraph_version = "unknown"
    return {
        "schema_version": SCHEMA_VERSION,
        "benchmark_version": BENCHMARK_VERSION,
        "success": True,
        "error": None,
        "profiled": False,
        "git_sha": git_sha,
        "python_version": platform.python_version(),
        "langgraph_version": langgraph_version,
        "platform": platform.platform(),
        "mode": case.mode,
        "backend": case.backend,
        "scenario": case.scenario,
        "snapshot_frequency": case.snapshot_frequency,
        "update_count": case.update_count,
        "payload_bytes": case.payload_bytes,
        "repetition": case.repetition,
        "seed": case.seed,
    }


def _collect_storage_stats(collector: Callable[[], dict[str, int]]) -> dict[str, Any]:
    """Keep timing data usable when a saver's diagnostic layout changes."""
    try:
        return {**collector(), "storage_stats_error": None}
    except Exception as exc:
        return {
            **dict.fromkeys(_STORAGE_STAT_FIELDS),
            "storage_stats_error": _safe_error(exc),
        }


def _memory_storage_stats(saver: InMemorySaver, thread_id: str) -> dict[str, int]:
    checkpoint_rows = 0
    checkpoint_bytes = 0
    for namespace in saver.storage.get(thread_id, {}).values():
        for checkpoint, metadata, _parent_id in namespace.values():
            checkpoint_rows += 1
            checkpoint_bytes += len(checkpoint[1]) + len(metadata[1])
    for (stored_thread_id, _namespace, _channel, _version), (_type_tag, blob) in saver.blobs.items():
        if stored_thread_id == thread_id:
            checkpoint_bytes += len(blob)

    write_rows = 0
    write_bytes = 0
    for (stored_thread_id, _namespace, _checkpoint_id), writes in saver.writes.items():
        if stored_thread_id != thread_id:
            continue
        for _task_id, _channel, (_type_tag, blob), _task_path in writes.values():
            write_rows += 1
            write_bytes += len(blob)
    return {
        "logical_checkpoint_bytes": checkpoint_bytes,
        "logical_write_bytes": write_bytes,
        "checkpoint_rows": checkpoint_rows,
        "write_rows": write_rows,
    }


def _sqlite_storage_stats(saver: SqliteSaver, thread_id: str) -> dict[str, int]:
    with saver.cursor(transaction=False) as cursor:
        checkpoint_rows, checkpoint_bytes = cursor.execute(
            "SELECT COUNT(*), COALESCE(SUM(length(checkpoint) + length(metadata)), 0) FROM checkpoints WHERE thread_id = ?",
            (thread_id,),
        ).fetchone()
        write_rows, write_bytes = cursor.execute(
            "SELECT COUNT(*), COALESCE(SUM(length(value)), 0) FROM writes WHERE thread_id = ?",
            (thread_id,),
        ).fetchone()
    return {
        "logical_checkpoint_bytes": int(checkpoint_bytes),
        "logical_write_bytes": int(write_bytes),
        "checkpoint_rows": int(checkpoint_rows),
        "write_rows": int(write_rows),
    }


def _write_and_read(case: BenchmarkCase, saver: Any, messages: list[BaseMessage]) -> tuple[dict[str, Any], list[AnyMessage]]:
    graph = _build_graph(case.mode, saver)
    accessor = CheckpointStateAccessor.bind(graph, saver, mode=case.mode)
    config = _config(case)
    update_latencies: list[float] = []
    write_start = time.perf_counter()
    for message in messages:
        update_start = time.perf_counter()
        graph.invoke({"messages": [message]}, config)
        update_latencies.append((time.perf_counter() - update_start) * 1000)
    write_total_ms = (time.perf_counter() - write_start) * 1000

    warm_start = time.perf_counter()
    snapshot = accessor.get(config)
    warm_read_ms = (time.perf_counter() - warm_start) * 1000
    materialized = list(snapshot.values.get("messages", []))
    metrics = {
        "write_total_ms": write_total_ms,
        "write_p50_ms": _percentile(update_latencies, 50),
        "write_p95_ms": _percentile(update_latencies, 95),
        "write_p99_ms": _percentile(update_latencies, 99),
        "write_first_window_ms": _window_median(update_latencies, "first"),
        "write_middle_window_ms": _window_median(update_latencies, "middle"),
        "write_last_window_ms": _window_median(update_latencies, "last"),
        "warm_read_ms": warm_read_ms,
    }
    return metrics, materialized


def _cold_read(case: BenchmarkCase, saver: Any) -> tuple[float, list[AnyMessage]]:
    graph = _build_graph(case.mode, saver)
    accessor = CheckpointStateAccessor.bind(graph, saver, mode=case.mode)
    gc.collect()
    start = time.perf_counter()
    snapshot = accessor.get(_config(case))
    elapsed_ms = (time.perf_counter() - start) * 1000
    return elapsed_ms, list(snapshot.values.get("messages", []))


def _validate_materialized(case: BenchmarkCase, expected: list[BaseMessage], warm: list[AnyMessage], cold: list[AnyMessage]) -> tuple[int, str]:
    expected_digest = _canonical_messages_digest(expected)
    warm_digest = _canonical_messages_digest(warm)
    cold_digest = _canonical_messages_digest(cold)
    if len(warm) != case.update_count or len(cold) != case.update_count:
        raise AssertionError(f"expected {case.update_count} messages, materialized warm={len(warm)} cold={len(cold)}")
    if warm_digest != expected_digest or cold_digest != expected_digest:
        raise AssertionError("materialized message content or ordering differs from deterministic input")
    return len(cold), cold_digest


def _run_memory_case(case: BenchmarkCase, messages: list[BaseMessage]) -> dict[str, Any]:
    saver = InMemorySaver()
    metrics, warm = _write_and_read(case, saver, messages)
    stats = _collect_storage_stats(lambda: _memory_storage_stats(saver, _config(case)["configurable"]["thread_id"]))
    cold_read_ms, cold = _cold_read(case, saver)
    actual_count, digest = _validate_materialized(case, messages, warm, cold)
    return {
        **metrics,
        **stats,
        "cold_read_ms": cold_read_ms,
        # InMemorySaver has no durable external storage to reopen. The cold
        # sample rebuilds the graph/channel table over the same saver only.
        "saver_reopen_ms": 0.0,
        "db_bytes": None,
        "wal_bytes": None,
        "shm_bytes": None,
        "durable_db_bytes": None,
        "expected_message_count": case.update_count,
        "actual_message_count": actual_count,
        "content_sha256": digest,
    }


def _run_sqlite_case(case: BenchmarkCase, messages: list[BaseMessage], db_path: Path) -> dict[str, Any]:
    with SqliteSaver.from_conn_string(str(db_path)) as saver:
        saver.setup()
        metrics, warm = _write_and_read(case, saver, messages)
        stats = _collect_storage_stats(lambda: _sqlite_storage_stats(saver, _config(case)["configurable"]["thread_id"]))
        db_bytes = _file_size(db_path)
        wal_bytes = _file_size(Path(f"{db_path}-wal"))
        shm_bytes = _file_size(Path(f"{db_path}-shm"))

    durable_db_bytes = _file_size(db_path)
    reopen_start = time.perf_counter()
    with SqliteSaver.from_conn_string(str(db_path)) as reopened:
        reopened.setup()
        saver_reopen_ms = (time.perf_counter() - reopen_start) * 1000
        cold_read_ms, cold = _cold_read(case, reopened)

    actual_count, digest = _validate_materialized(case, messages, warm, cold)
    return {
        **metrics,
        **stats,
        "cold_read_ms": cold_read_ms,
        "saver_reopen_ms": saver_reopen_ms,
        "db_bytes": db_bytes,
        "wal_bytes": wal_bytes,
        "shm_bytes": shm_bytes,
        "durable_db_bytes": durable_db_bytes,
        "expected_message_count": case.update_count,
        "actual_message_count": actual_count,
        "content_sha256": digest,
    }


def _run_case(case: BenchmarkCase, *, work_dir: Path) -> dict[str, Any]:
    row = _base_row(case)
    messages = [_message_for_update(index, case.payload_bytes) for index in range(case.update_count)]
    try:
        if case.backend == "memory":
            measured = _run_memory_case(case, messages)
        else:
            measured = _run_sqlite_case(case, messages, work_dir / "checkpoint-benchmark.sqlite")

        reducer_writes = [[message] for message in messages]
        reducer_start = time.perf_counter()
        reduced = merge_message_writes([], reducer_writes)
        reducer_replay_ms = (time.perf_counter() - reducer_start) * 1000
        if _canonical_messages_digest(reduced) != _canonical_messages_digest(messages):
            raise AssertionError("standalone reducer diagnostic produced incorrect state")

        row.update(measured)
        row["reducer_replay_ms"] = reducer_replay_ms
        row["peak_rss_bytes"] = _peak_rss_bytes()
    except Exception as exc:
        row["success"] = False
        row["error"] = _safe_error(exc, work_dir=work_dir)
    return row


def _run_profiled_case(case: BenchmarkCase, *, work_dir: Path, profile_path: Path) -> dict[str, Any]:
    return _common.run_profiled(_run_case, case, work_dir=work_dir, profile_path=profile_path)


def _comparison_key(row: dict[str, Any]) -> tuple[Any, ...]:
    return tuple(
        row.get(field)
        for field in (
            "backend",
            "scenario",
            "snapshot_frequency",
            "update_count",
            "payload_bytes",
            "repetition",
        )
    )


def _validate_cross_mode_rows(rows: list[dict[str, Any]]) -> None:
    grouped: dict[tuple[Any, ...], list[dict[str, Any]]] = {}
    for row in rows:
        grouped.setdefault(_comparison_key(row), []).append(row)
    for group in grouped.values():
        successful = [row for row in group if row.get("success")]
        modes = {row.get("mode") for row in successful}
        if not {"full", "delta"}.issubset(modes):
            continue
        signatures = {(row.get("actual_message_count"), row.get("content_sha256")) for row in successful if row.get("mode") in {"full", "delta"}}
        if len(signatures) == 1:
            continue
        for row in successful:
            row["success"] = False
            row["error"] = "cross-mode materialized state mismatch"


def _failure_row(case: BenchmarkCase, error: str) -> dict[str, Any]:
    row = _base_row(case)
    row["success"] = False
    row["error"] = _safe_error(error)
    return row


def _profile_filename(case: BenchmarkCase) -> str:
    return f"{case.backend}-{case.mode}-updates-{case.update_count}-payload-{case.payload_bytes}-rep-{case.repetition}.prof"


def _run_child_case(case: BenchmarkCase, *, timeout_seconds: float, git_sha: str, profile_dir: Path | None = None) -> dict[str, Any]:
    worker_args = ["--worker-case", json.dumps(asdict(case), separators=(",", ":"))]
    if profile_dir is not None:
        worker_args.extend(["--worker-profile", str(profile_dir / _profile_filename(case))])
    return _common.run_child_case(
        script=Path(__file__).resolve(),
        worker_args=worker_args,
        failure_row=lambda error: _failure_row(case, error),
        timeout_seconds=timeout_seconds,
        git_sha=git_sha,
    )


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Benchmark full and delta checkpoint message channels")
    parser.add_argument("--modes", default="full,delta", help="Comma-separated modes (default: full,delta)")
    parser.add_argument("--backends", default="memory,sqlite", help="Comma-separated backends (default: memory,sqlite)")
    parser.add_argument("--updates", default="10,100", help="Comma-separated message update counts (default: 10,100)")
    parser.add_argument("--payload-bytes", default="128", help="Comma-separated exact message content sizes (default: 128)")
    parser.add_argument("--repetitions", type=int, default=3)
    parser.add_argument("--seed", type=int, default=1)
    parser.add_argument("--timeout-seconds", type=float, default=900)
    parser.add_argument(
        "--max-estimated-full-bytes",
        type=int,
        default=DEFAULT_MAX_ESTIMATED_FULL_BYTES,
        help=("Skip comparable pairs whose estimated cumulative full-mode payload exceeds this value. The cap applies only when full mode is selected; delta-only diagnostics bypass it."),
    )
    parser.add_argument("--allow-large-cases", action="store_true", help="Disable the estimated cumulative full-payload safety cap")
    parser.add_argument("--profile-dir", type=Path, help="Write one cProfile file per case; profiling inflates timings")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--worker-case", help=argparse.SUPPRESS)
    parser.add_argument("--worker-profile", type=Path, help=argparse.SUPPRESS)
    return parser


def _worker_main(encoded_case: str, *, profile_path: Path | None = None) -> int:
    try:
        case = BenchmarkCase(**json.loads(encoded_case))
    except (TypeError, ValueError, json.JSONDecodeError) as exc:
        print(json.dumps({"schema_version": SCHEMA_VERSION, "benchmark_version": BENCHMARK_VERSION, "success": False, "error": _safe_error(exc)}, separators=(",", ":")))
        return 2
    with tempfile.TemporaryDirectory(prefix="deerflow-checkpoint-benchmark-") as temp_dir:
        if profile_path is None:
            row = _run_case(case, work_dir=Path(temp_dir))
        else:
            row = _run_profiled_case(case, work_dir=Path(temp_dir), profile_path=profile_path)
    print(json.dumps(row, ensure_ascii=False, separators=(",", ":")))
    return 0 if row.get("success") else 1


def main(argv: list[str] | None = None) -> int:
    parser = _build_parser()
    args = parser.parse_args(argv)
    if args.worker_case is not None:
        return _worker_main(args.worker_case, profile_path=args.worker_profile)
    if args.output is None:
        parser.error("--output is required")
    try:
        modes = _parse_choice_csv(args.modes, option="--modes", choices=_MODES)
        backends = _parse_choice_csv(args.backends, option="--backends", choices=_BACKENDS)
        updates = _parse_positive_int_csv(args.updates, option="--updates")
        payload_bytes = _parse_positive_int_csv(args.payload_bytes, option="--payload-bytes")
        if args.repetitions <= 0:
            raise ValueError("--repetitions must be positive")
        if args.timeout_seconds <= 0:
            raise ValueError("--timeout-seconds must be positive")
        if not args.allow_large_cases and args.max_estimated_full_bytes <= 0:
            raise ValueError("--max-estimated-full-bytes must be positive")
    except ValueError as exc:
        parser.error(str(exc))

    cases = _expand_cases(
        modes=modes,
        backends=backends,
        update_counts=updates,
        payload_bytes=payload_bytes,
        repetitions=args.repetitions,
        seed=args.seed,
    )
    cases, skipped = _filter_oversized_pairs(cases, max_bytes=None if args.allow_large_cases else args.max_estimated_full_bytes)
    if skipped:
        skipped_pairs = len(skipped) // max(1, len(modes))
        print(
            f"Skipping {skipped_pairs} oversized comparable case pair(s); use --allow-large-cases to run them.",
            file=sys.stderr,
        )
    if not cases:
        print("No benchmark cases remain after applying the safety cap.", file=sys.stderr)
        return 2

    git_sha = _resolve_git_sha()
    rows: list[dict[str, Any]] = []
    for index, case in enumerate(cases, start=1):
        print(
            f"[{index}/{len(cases)}] {case.backend} {case.mode} updates={case.update_count} payload={case.payload_bytes} repetition={case.repetition}",
            file=sys.stderr,
        )
        rows.append(_run_child_case(case, timeout_seconds=args.timeout_seconds, git_sha=git_sha, profile_dir=args.profile_dir))
    _validate_cross_mode_rows(rows)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    with args.output.open("w", encoding="utf-8") as output_file:
        for row in rows:
            output_file.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
    failures = sum(1 for row in rows if not row.get("success"))
    print(f"Wrote {len(rows)} result row(s) to {args.output}; failures={failures}", file=sys.stderr)
    return 1 if failures else 0


if __name__ == "__main__":
    raise SystemExit(main())