scripts/pixelmodel_robustness/storage.py

"""Atomic and append-only study storage."""

from __future__ import annotations

import hashlib
import json
import os
import tempfile
from pathlib import Path
from typing import Any


def atomic_json(path: str | Path, payload: dict[str, Any]) -> None:
    """同一ディレクトリ内のtemporary fileからmanifestをatomic replaceする。"""
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary = tempfile.mkstemp(prefix=f".{target.name}.", dir=target.parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as stream:
            json.dump(payload, stream, ensure_ascii=False, indent=2, sort_keys=True)
            stream.write("\n")
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, target)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def load_manifest(path: str | Path) -> dict[str, Any]:
    """存在するmanifestを読み、無ければ空のpartial stateを返す。"""
    target = Path(path)
    return json.loads(target.read_text(encoding="utf-8")) if target.exists() else {"status": "partial"}


def save_jsonl_row(path: str | Path, row: dict[str, Any]) -> None:
    """canonical rowをidentity upsertし、atomicに保存する。"""
    target = Path(path)
    target.parent.mkdir(parents=True, exist_ok=True)
    rows = read_jsonl_rows(target)
    identity = _row_identity(row)
    replaced = False
    for index, existing in enumerate(rows):
        if _row_identity(existing) == identity:
            rows[index] = row
            replaced = True
            break
    if not replaced:
        rows.append(row)
    fd, temporary = tempfile.mkstemp(prefix=f".{target.name}.", dir=target.parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as stream:
            for item in rows:
                stream.write(json.dumps(item, ensure_ascii=False, allow_nan=False) + "\n")
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(temporary, target)
    finally:
        if os.path.exists(temporary):
            os.unlink(temporary)


def _row_identity(row: dict[str, Any]) -> str:
    """canonical raw rowの重複排除キー。"""
    fields = ("condition", "prompt_id", "stage", "row_type")
    return json.dumps({key: row.get(key) for key in fields}, sort_keys=True)


def read_jsonl_rows(path: str | Path) -> list[dict[str, Any]]:
    """canonical raw JSONLを読み、identity keyでdedupeする。"""
    target = Path(path)
    if not target.exists():
        return []
    rows: dict[str, dict[str, Any]] = {}
    for line in target.read_text(encoding="utf-8").splitlines():
        if line.strip():
            row = json.loads(line)
            rows[_row_identity(row)] = row
    return list(rows.values())


def assert_resume_identity(existing: dict[str, Any], current: dict[str, Any]) -> None:
    """scientific identityが違うresumeをfail-closedで拒否する。"""
    if existing == {"status": "partial"}:
        return
    keys = ("source_sha256", "config_sha256", "prompt_sha256", "model_sha256", "real_content_hash", "cli_identity")
    mismatches = [key for key in keys if key not in existing or existing.get(key) != current.get(key)]
    if mismatches:
        raise ValueError(f"resume identity mismatch: {', '.join(mismatches)}")


def validate_real_manifest(path: str | Path) -> dict[str, Any]:
    """FID用real manifestのpath/hashを検証する。"""
    target = Path(path)
    payload = json.loads(target.read_text(encoding="utf-8"))
    images = payload.get("images") if isinstance(payload, dict) else payload
    if not isinstance(images, list) or not images:
        raise ValueError("real images manifest must contain a non-empty images list")
    enriched: list[dict[str, Any]] = []
    ordered = hashlib.sha256()
    for item in images:
        image = Path(item["path"] if isinstance(item, dict) else item)
        if not image.is_absolute():
            image = target.parent / image
        if not image.exists():
            raise FileNotFoundError(image)
        digest = hashlib.sha256()
        with image.open("rb") as stream:
            for chunk in iter(lambda: stream.read(1024 * 1024), b""):
                digest.update(chunk)
        item_data = dict(item) if isinstance(item, dict) else {"path": str(item)}
        item_data["sha256"] = digest.hexdigest()
        enriched.append(item_data)
        ordered.update(str(item_data["path"]).encode())
        ordered.update(b"\0")
        ordered.update(digest.digest())
    return {"path": str(target), "n": len(enriched), "manifest_sha256": hashlib.sha256(target.read_bytes()).hexdigest(), "real_content_hash": ordered.hexdigest(), "images": enriched}