mirror of
https://github.com/norandom/OpenUltraSAST.git
synced 2026-11-13 01:48:17 +00:00
310 lines
13 KiB
Python
310 lines
13 KiB
Python
|
"""Strict benchmark-owned persistence for canonical v2 artifacts."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import fcntl
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import tempfile
|
||
|
|
from collections.abc import Iterator, Mapping, Sequence
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import TypeAlias, cast
|
||
|
|
|
||
|
|
from openultrasast.core.canonical import JsonValue, content_digest
|
||
|
|
|
||
|
|
BENCHMARK_ARTIFACT_SCHEMA = "openultrasast.benchmark-artifact.v2"
|
||
|
|
ArtifactPayload: TypeAlias = Mapping[str, JsonValue]
|
||
|
|
ArtifactRows: TypeAlias = Sequence[Mapping[str, JsonValue]]
|
||
|
|
_RESERVED = frozenset({"schema_version", "artifact_kind", "artifact_digest"})
|
||
|
|
_SEQUENCE_KEYS = frozenset({"artifact_mode", "sequence_count", "sequence_digest", "rows"})
|
||
|
|
_JSONL_HEADER_KEYS = frozenset({"artifact_mode", "sequence_count", "sequence_digest"})
|
||
|
|
_JSONL_ROW_KEYS = frozenset({"sequence_index", "sequence_count", "sequence_digest", "artifact"})
|
||
|
|
|
||
|
|
|
||
|
|
def _canonical_value(value: object) -> JsonValue:
|
||
|
|
if value is None or isinstance(value, (str, bool, int, float)):
|
||
|
|
return value
|
||
|
|
if isinstance(value, Mapping):
|
||
|
|
return {str(key): _canonical_value(item) for key, item in value.items()}
|
||
|
|
if isinstance(value, Sequence) and not isinstance(value, (str, bytes)):
|
||
|
|
return tuple(_canonical_value(item) for item in value)
|
||
|
|
raise ValueError(f"unsupported benchmark artifact value: {type(value).__name__}")
|
||
|
|
|
||
|
|
|
||
|
|
def _kind(kind: str) -> str:
|
||
|
|
if not isinstance(kind, str) or not kind:
|
||
|
|
raise ValueError("benchmark artifact kind must be nonempty text")
|
||
|
|
return kind
|
||
|
|
|
||
|
|
|
||
|
|
def _artifact_digest(kind: str, payload: Mapping[str, JsonValue]) -> str:
|
||
|
|
return content_digest(
|
||
|
|
{
|
||
|
|
"schema_version": BENCHMARK_ARTIFACT_SCHEMA,
|
||
|
|
"artifact_kind": kind,
|
||
|
|
"payload": _canonical_value(payload),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def benchmark_artifact(kind: str, payload: ArtifactPayload) -> dict[str, JsonValue]:
|
||
|
|
"""Create one v2 artifact from an unwrapped payload."""
|
||
|
|
kind = _kind(kind)
|
||
|
|
if not isinstance(payload, Mapping):
|
||
|
|
raise ValueError("benchmark artifact payload must be a mapping")
|
||
|
|
if _RESERVED.intersection(payload):
|
||
|
|
raise ValueError("benchmark artifact payload uses reserved metadata")
|
||
|
|
value = {str(key): _canonical_value(item) for key, item in payload.items()}
|
||
|
|
return {
|
||
|
|
**value,
|
||
|
|
"schema_version": BENCHMARK_ARTIFACT_SCHEMA,
|
||
|
|
"artifact_kind": kind,
|
||
|
|
"artifact_digest": _artifact_digest(kind, value),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def benchmark_payload(value: object, *, expected_kind: str) -> dict[str, JsonValue]:
|
||
|
|
"""Strictly validate one v2 artifact and return its domain payload."""
|
||
|
|
expected_kind = _kind(expected_kind)
|
||
|
|
if not isinstance(value, dict):
|
||
|
|
raise ValueError("benchmark artifact metadata is missing")
|
||
|
|
if _RESERVED.intersection(value) != _RESERVED:
|
||
|
|
raise ValueError("benchmark artifact metadata is missing or incomplete")
|
||
|
|
raw = {str(key): item for key, item in value.items() if key not in _RESERVED}
|
||
|
|
canonical_raw = cast(dict[str, JsonValue], _canonical_value(raw))
|
||
|
|
if (
|
||
|
|
value["schema_version"] != BENCHMARK_ARTIFACT_SCHEMA
|
||
|
|
or value["artifact_kind"] != expected_kind
|
||
|
|
or value["artifact_digest"] != _artifact_digest(expected_kind, canonical_raw)
|
||
|
|
):
|
||
|
|
raise ValueError("benchmark artifact metadata does not match payload")
|
||
|
|
return cast(dict[str, JsonValue], raw)
|
||
|
|
|
||
|
|
|
||
|
|
def encode_benchmark_artifact(kind: str, payload: ArtifactPayload, *, indent: int | None = 2) -> str:
|
||
|
|
"""Encode one canonical v2 artifact."""
|
||
|
|
return json.dumps(benchmark_artifact(kind, payload), sort_keys=True, indent=indent, ensure_ascii=False)
|
||
|
|
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def _target_lock(path: Path) -> Iterator[None]:
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
descriptor = os.open(path.with_name(f".{path.name}.lock"), os.O_CREAT | os.O_RDWR, 0o600)
|
||
|
|
try:
|
||
|
|
fcntl.flock(descriptor, fcntl.LOCK_EX)
|
||
|
|
yield
|
||
|
|
finally:
|
||
|
|
fcntl.flock(descriptor, fcntl.LOCK_UN)
|
||
|
|
os.close(descriptor)
|
||
|
|
|
||
|
|
|
||
|
|
def _write_locked(path: Path, content: bytes, *, replace: bool) -> None:
|
||
|
|
if path.exists():
|
||
|
|
existing = path.read_bytes()
|
||
|
|
if existing == content:
|
||
|
|
return
|
||
|
|
if not replace:
|
||
|
|
raise ValueError(f"benchmark artifact conflict: {path}")
|
||
|
|
temporary_name: str | None = None
|
||
|
|
try:
|
||
|
|
with tempfile.NamedTemporaryFile(mode="wb", dir=path.parent, prefix=f".{path.name}.", suffix=".tmp", delete=False) as temporary:
|
||
|
|
temporary_name = temporary.name
|
||
|
|
temporary.write(content)
|
||
|
|
temporary.flush()
|
||
|
|
os.fsync(temporary.fileno())
|
||
|
|
os.replace(temporary_name, path)
|
||
|
|
temporary_name = None
|
||
|
|
finally:
|
||
|
|
if temporary_name is not None:
|
||
|
|
Path(temporary_name).unlink(missing_ok=True)
|
||
|
|
|
||
|
|
|
||
|
|
def _write_bytes(path: Path, content: bytes, *, replace: bool) -> None:
|
||
|
|
path = Path(path)
|
||
|
|
with _target_lock(path):
|
||
|
|
_write_locked(path, content, replace=replace)
|
||
|
|
|
||
|
|
|
||
|
|
def write_benchmark_artifact(path: Path, kind: str, payload: ArtifactPayload, *, replace: bool = False) -> None:
|
||
|
|
"""Atomically create one artifact; identical writes are idempotent."""
|
||
|
|
_write_bytes(path, (encode_benchmark_artifact(kind, payload) + "\n").encode(), replace=replace)
|
||
|
|
|
||
|
|
|
||
|
|
def read_benchmark_artifact(path: Path, expected_kind: str) -> dict[str, JsonValue]:
|
||
|
|
"""Read and strictly validate one artifact from disk."""
|
||
|
|
try:
|
||
|
|
value = json.loads(Path(path).read_text())
|
||
|
|
except (OSError, json.JSONDecodeError) as error:
|
||
|
|
raise ValueError(f"benchmark artifact unreadable: {path}") from error
|
||
|
|
return benchmark_payload(value, expected_kind=expected_kind)
|
||
|
|
|
||
|
|
|
||
|
|
def _ordered_digest(artifacts: Sequence[Mapping[str, JsonValue]]) -> str:
|
||
|
|
return content_digest(
|
||
|
|
tuple(
|
||
|
|
{"sequence_index": index, "artifact_digest": cast(str, artifact["artifact_digest"])} for index, artifact in enumerate(artifacts)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def _sequence_artifact(kind: str, rows: ArtifactRows) -> dict[str, JsonValue]:
|
||
|
|
artifacts = tuple(benchmark_artifact(f"{kind}.item", row) for row in rows)
|
||
|
|
members: tuple[JsonValue, ...] = tuple({"sequence_index": index, "artifact": artifact} for index, artifact in enumerate(artifacts))
|
||
|
|
return benchmark_artifact(
|
||
|
|
kind,
|
||
|
|
{
|
||
|
|
"artifact_mode": "ordered-sequence",
|
||
|
|
"sequence_count": len(artifacts),
|
||
|
|
"sequence_digest": _ordered_digest(artifacts),
|
||
|
|
"rows": members,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def write_benchmark_sequence(path: Path, kind: str, rows: ArtifactRows, *, replace: bool = False) -> None:
|
||
|
|
"""Persist an exact ordered collection with count, indices, and aggregate digest."""
|
||
|
|
value = _sequence_artifact(_kind(kind), rows)
|
||
|
|
content = (json.dumps(value, sort_keys=True, indent=2, ensure_ascii=False) + "\n").encode()
|
||
|
|
_write_bytes(path, content, replace=replace)
|
||
|
|
|
||
|
|
|
||
|
|
def _sequence_payload(value: object, kind: str) -> tuple[dict[str, JsonValue], ...]:
|
||
|
|
payload = benchmark_payload(value, expected_kind=kind)
|
||
|
|
if set(payload) != _SEQUENCE_KEYS or payload.get("artifact_mode") != "ordered-sequence":
|
||
|
|
raise ValueError("benchmark sequence metadata does not match contract")
|
||
|
|
count = payload.get("sequence_count")
|
||
|
|
sequence_digest = payload.get("sequence_digest")
|
||
|
|
members = payload.get("rows")
|
||
|
|
if (
|
||
|
|
not isinstance(count, int)
|
||
|
|
or count < 0
|
||
|
|
or not isinstance(sequence_digest, str)
|
||
|
|
or not isinstance(members, Sequence)
|
||
|
|
or isinstance(members, (str, bytes))
|
||
|
|
):
|
||
|
|
raise ValueError("benchmark sequence metadata is invalid")
|
||
|
|
if len(members) != count:
|
||
|
|
raise ValueError("benchmark sequence count mismatch")
|
||
|
|
rows: list[dict[str, JsonValue]] = []
|
||
|
|
artifacts: list[dict[str, JsonValue]] = []
|
||
|
|
for index, member in enumerate(members):
|
||
|
|
if not isinstance(member, dict) or set(member) != {"sequence_index", "artifact"} or member.get("sequence_index") != index:
|
||
|
|
raise ValueError("benchmark sequence index mismatch")
|
||
|
|
artifact = member.get("artifact")
|
||
|
|
rows.append(benchmark_payload(artifact, expected_kind=f"{kind}.item"))
|
||
|
|
artifacts.append(cast(dict[str, JsonValue], artifact))
|
||
|
|
digests = tuple(cast(str, artifact["artifact_digest"]) for artifact in artifacts)
|
||
|
|
if len(set(digests)) != len(digests) or sequence_digest != _ordered_digest(artifacts):
|
||
|
|
raise ValueError("benchmark sequence digest mismatch or duplicate row")
|
||
|
|
return tuple(rows)
|
||
|
|
|
||
|
|
|
||
|
|
def read_benchmark_sequence(path: Path, expected_kind: str) -> tuple[dict[str, JsonValue], ...]:
|
||
|
|
"""Read and validate a complete ordered collection."""
|
||
|
|
try:
|
||
|
|
value = json.loads(Path(path).read_text())
|
||
|
|
except (OSError, json.JSONDecodeError) as error:
|
||
|
|
raise ValueError(f"benchmark sequence unreadable: {path}") from error
|
||
|
|
return _sequence_payload(value, _kind(expected_kind))
|
||
|
|
|
||
|
|
|
||
|
|
def _jsonl_bytes(kind: str, rows: ArtifactRows) -> bytes:
|
||
|
|
items = tuple(benchmark_artifact(f"{kind}.item", row) for row in rows)
|
||
|
|
sequence_digest = _ordered_digest(items)
|
||
|
|
header = benchmark_artifact(
|
||
|
|
f"{kind}.sequence",
|
||
|
|
{"artifact_mode": "ordered-jsonl", "sequence_count": len(items), "sequence_digest": sequence_digest},
|
||
|
|
)
|
||
|
|
encoded = [json.dumps(header, sort_keys=True, ensure_ascii=False)]
|
||
|
|
for index, artifact in enumerate(items):
|
||
|
|
encoded.append(
|
||
|
|
json.dumps(
|
||
|
|
benchmark_artifact(
|
||
|
|
f"{kind}.row",
|
||
|
|
{
|
||
|
|
"sequence_index": index,
|
||
|
|
"sequence_count": len(items),
|
||
|
|
"sequence_digest": sequence_digest,
|
||
|
|
"artifact": artifact,
|
||
|
|
},
|
||
|
|
),
|
||
|
|
sort_keys=True,
|
||
|
|
ensure_ascii=False,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return ("\n".join(encoded) + "\n").encode()
|
||
|
|
|
||
|
|
|
||
|
|
def _jsonl_rows(content: bytes, kind: str) -> tuple[dict[str, JsonValue], ...]:
|
||
|
|
try:
|
||
|
|
values = tuple(json.loads(line) for line in content.decode().splitlines())
|
||
|
|
except (UnicodeDecodeError, json.JSONDecodeError) as error:
|
||
|
|
raise ValueError("benchmark JSONL is unreadable") from error
|
||
|
|
if not values:
|
||
|
|
raise ValueError("benchmark JSONL metadata is missing")
|
||
|
|
header = benchmark_payload(values[0], expected_kind=f"{kind}.sequence")
|
||
|
|
if set(header) != _JSONL_HEADER_KEYS or header.get("artifact_mode") != "ordered-jsonl":
|
||
|
|
raise ValueError("benchmark JSONL metadata does not match contract")
|
||
|
|
count = header.get("sequence_count")
|
||
|
|
sequence_digest = header.get("sequence_digest")
|
||
|
|
if not isinstance(count, int) or count < 0 or not isinstance(sequence_digest, str) or len(values) != count + 1:
|
||
|
|
raise ValueError("benchmark JSONL count mismatch")
|
||
|
|
rows: list[dict[str, JsonValue]] = []
|
||
|
|
items: list[dict[str, JsonValue]] = []
|
||
|
|
for index, value in enumerate(values[1:]):
|
||
|
|
member = benchmark_payload(value, expected_kind=f"{kind}.row")
|
||
|
|
if set(member) != _JSONL_ROW_KEYS:
|
||
|
|
raise ValueError("benchmark JSONL row metadata does not match contract")
|
||
|
|
if (
|
||
|
|
member.get("sequence_index") != index
|
||
|
|
or member.get("sequence_count") != count
|
||
|
|
or member.get("sequence_digest") != sequence_digest
|
||
|
|
):
|
||
|
|
raise ValueError("benchmark JSONL sequence binding mismatch")
|
||
|
|
artifact = member.get("artifact")
|
||
|
|
rows.append(benchmark_payload(artifact, expected_kind=f"{kind}.item"))
|
||
|
|
items.append(cast(dict[str, JsonValue], artifact))
|
||
|
|
digests = tuple(cast(str, item["artifact_digest"]) for item in items)
|
||
|
|
if len(set(digests)) != len(digests) or sequence_digest != _ordered_digest(items):
|
||
|
|
raise ValueError("benchmark JSONL digest mismatch or duplicate row")
|
||
|
|
return tuple(rows)
|
||
|
|
|
||
|
|
|
||
|
|
def write_benchmark_jsonl(path: Path, kind: str, rows: ArtifactRows, *, replace: bool = False) -> None:
|
||
|
|
"""Atomically persist an ordered JSONL manifest and its bound rows."""
|
||
|
|
_write_bytes(path, _jsonl_bytes(_kind(kind), rows), replace=replace)
|
||
|
|
|
||
|
|
|
||
|
|
def read_benchmark_jsonl(path: Path, expected_kind: str) -> tuple[dict[str, JsonValue], ...]:
|
||
|
|
"""Read and validate a complete ordered JSONL artifact."""
|
||
|
|
try:
|
||
|
|
content = Path(path).read_bytes()
|
||
|
|
except OSError as error:
|
||
|
|
raise ValueError(f"benchmark JSONL unreadable: {path}") from error
|
||
|
|
return _jsonl_rows(content, _kind(expected_kind))
|
||
|
|
|
||
|
|
|
||
|
|
def append_benchmark_jsonl(path: Path, kind: str, row: ArtifactPayload) -> None:
|
||
|
|
"""Append one row under the target lock and rewrite the bound sequence atomically."""
|
||
|
|
path = Path(path)
|
||
|
|
kind = _kind(kind)
|
||
|
|
with _target_lock(path):
|
||
|
|
rows = _jsonl_rows(path.read_bytes(), kind) if path.exists() else ()
|
||
|
|
_write_locked(path, _jsonl_bytes(kind, (*rows, row)), replace=True)
|
||
|
|
|
||
|
|
|
||
|
|
__all__ = [
|
||
|
|
"BENCHMARK_ARTIFACT_SCHEMA",
|
||
|
|
"append_benchmark_jsonl",
|
||
|
|
"benchmark_artifact",
|
||
|
|
"benchmark_payload",
|
||
|
|
"encode_benchmark_artifact",
|
||
|
|
"read_benchmark_artifact",
|
||
|
|
"read_benchmark_jsonl",
|
||
|
|
"read_benchmark_sequence",
|
||
|
|
"write_benchmark_artifact",
|
||
|
|
"write_benchmark_jsonl",
|
||
|
|
"write_benchmark_sequence",
|
||
|
|
]
|