neuron-tai/packages/node/meshnet_node/performance_contract.py

"""Versioned performance contract metadata for DGR-001.

This module intentionally captures the *contract* first: the model family,
architecture alignment, benchmark lanes, and stop condition that later benchmark
runs must satisfy. It does not download or execute a model.
"""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from pathlib import Path

SCHEMA_VERSION = 1
CONTRACT_ID = "DGR-001"
DEFAULT_OUTPUT_PATH = Path(".scratch/distributed-gguf-runtime/evidence/DGR-001/performance-contract.json")


@dataclass(frozen=True)
class ModelTarget:
    """Architecture-aligned model target for the DGR-001 benchmark contract."""

    name: str
    architecture: str
    safetensors_repo: str
    gguf_repo: str
    gguf_quant: str
    gguf_size_gb: float
    rationale: str

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "architecture": self.architecture,
            "safetensors_repo": self.safetensors_repo,
            "gguf_repo": self.gguf_repo,
            "gguf_quant": self.gguf_quant,
            "gguf_size_gb": self.gguf_size_gb,
            "rationale": self.rationale,
        }


@dataclass(frozen=True)
class BenchmarkLane:
    """One side of the comparison the contract requires."""

    id: str
    runtime: str
    recipe: str
    concurrency_levels: tuple[int, ...]

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "runtime": self.runtime,
            "recipe": self.recipe,
            "concurrency_levels": list(self.concurrency_levels),
        }


@dataclass(frozen=True)
class PerformanceContract:
    """Machine-readable contract for the DGR-001 benchmark story."""

    schema_version: int
    story_id: str
    model_target: ModelTarget
    benchmark_lanes: tuple[BenchmarkLane, ...]
    metrics: tuple[str, ...]
    stop_condition: str
    notes: tuple[str, ...] = ()

    def to_dict(self) -> dict:
        return {
            "schema_version": self.schema_version,
            "story_id": self.story_id,
            "model_target": self.model_target.to_dict(),
            "benchmark_lanes": [lane.to_dict() for lane in self.benchmark_lanes],
            "metrics": list(self.metrics),
            "stop_condition": self.stop_condition,
            "notes": list(self.notes),
        }

    def write_json(self, path: str | Path) -> Path:
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n", encoding="utf-8")
        return path


DEFAULT_CONTRACT = PerformanceContract(
    schema_version=SCHEMA_VERSION,
    story_id=CONTRACT_ID,
    model_target=ModelTarget(
        name="DeepSeek-V2-Lite-Chat",
        architecture="deepseek2",
        safetensors_repo="deepseek-ai/DeepSeek-V2-Lite-Chat",
        gguf_repo="second-state/DeepSeek-V2-Lite-Chat-GGUF",
        gguf_quant="Q2_K",
        gguf_size_gb=6.43,
        rationale=(
            "Smallest DeepSeek-family benchmark anchor that still points toward "
            "DeepSeek-V4-Flash; keeps the runtime on the DeepSeek2 path instead "
            "of falling back to a tiny but architecture-mismatched smoke model."
        ),
    ),
    benchmark_lanes=(
        BenchmarkLane(
            id="transformers-safetensors",
            runtime="transformers",
            recipe="current safetensors recipe",
            concurrency_levels=(1, 4),
        ),
        BenchmarkLane(
            id="llama-cpp-gguf",
            runtime="llama.cpp",
            recipe="whole-model GGUF recipe",
            concurrency_levels=(1, 4),
        ),
    ),
    metrics=(
        "ttft_ms",
        "prefill_tok_per_sec",
        "decode_tok_per_sec",
        "p50_latency_ms",
        "p95_latency_ms",
        "aggregate_throughput_tok_per_sec",
        "rss_bytes",
        "vram_bytes",
        "artifact_bytes",
        "failure_count",
        "output_drift",
    ),
    stop_condition=(
        "Stop if GGUF does not provide a meaningful speed or fit benefit over the "
        "safetensors baseline for the chosen DeepSeek-family model target."
    ),
    notes=(
        "Real model execution stays opt-in and must keep model artifacts on the mounted drive.",
        "Use the tiny fallback only for loader plumbing smoke tests; it does not replace the architecture-aligned baseline.",
    ),
)


def build_default_contract() -> PerformanceContract:
    return DEFAULT_CONTRACT


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Write the DGR-001 performance contract JSON")
    parser.add_argument("--json-out", type=Path, default=DEFAULT_OUTPUT_PATH, help="output JSON path")
    args = parser.parse_args(argv)
    contract = build_default_contract()
    path = contract.write_json(args.json_out)
    print(path)
    return 0


if __name__ == "__main__":  # pragma: no cover - CLI entry point
    raise SystemExit(main())