210 lines
7.2 KiB
Python
210 lines
7.2 KiB
Python
|
"""Atomic evidence writers using only the standard library and NumPy."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import tempfile
|
||
|
|
import uuid
|
||
|
|
from dataclasses import dataclass, field
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Mapping, Sequence
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from .hashing import file_sha256, to_jsonable
|
||
|
|
from .schema import SCHEMA_VERSION, STORAGE_FORMATS, ContractError
|
||
|
|
|
||
|
|
|
||
|
|
def utc_now() -> str:
|
||
|
|
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
||
|
|
|
||
|
|
|
||
|
|
def _flush_and_sync(stream) -> None:
|
||
|
|
stream.flush()
|
||
|
|
os.fsync(stream.fileno())
|
||
|
|
|
||
|
|
|
||
|
|
def atomic_write_json(path: Path, document: Mapping[str, Any]) -> None:
|
||
|
|
path = Path(path)
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
fd, temporary_name = tempfile.mkstemp(
|
||
|
|
prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
|
||
|
|
)
|
||
|
|
temporary = Path(temporary_name)
|
||
|
|
try:
|
||
|
|
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
||
|
|
json.dump(
|
||
|
|
to_jsonable(document),
|
||
|
|
stream,
|
||
|
|
indent=2,
|
||
|
|
sort_keys=True,
|
||
|
|
ensure_ascii=False,
|
||
|
|
allow_nan=False,
|
||
|
|
)
|
||
|
|
stream.write("\n")
|
||
|
|
_flush_and_sync(stream)
|
||
|
|
os.replace(temporary, path)
|
||
|
|
except Exception:
|
||
|
|
temporary.unlink(missing_ok=True)
|
||
|
|
raise
|
||
|
|
|
||
|
|
|
||
|
|
def atomic_write_jsonl(path: Path, rows: Sequence[Mapping[str, Any]]) -> None:
|
||
|
|
path = Path(path)
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
fd, temporary_name = tempfile.mkstemp(
|
||
|
|
prefix=f".{path.name}.", suffix=".tmp", dir=path.parent
|
||
|
|
)
|
||
|
|
temporary = Path(temporary_name)
|
||
|
|
try:
|
||
|
|
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
||
|
|
for row in rows:
|
||
|
|
stream.write(
|
||
|
|
json.dumps(
|
||
|
|
to_jsonable(row),
|
||
|
|
sort_keys=True,
|
||
|
|
separators=(",", ":"),
|
||
|
|
ensure_ascii=False,
|
||
|
|
allow_nan=False,
|
||
|
|
)
|
||
|
|
)
|
||
|
|
stream.write("\n")
|
||
|
|
_flush_and_sync(stream)
|
||
|
|
os.replace(temporary, path)
|
||
|
|
except Exception:
|
||
|
|
temporary.unlink(missing_ok=True)
|
||
|
|
raise
|
||
|
|
|
||
|
|
|
||
|
|
def atomic_write_npz(path: Path, arrays: Mapping[str, np.ndarray]) -> None:
|
||
|
|
path = Path(path)
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
temporary = path.parent / f".{path.stem}.{uuid.uuid4().hex}.tmp.npz"
|
||
|
|
try:
|
||
|
|
np.savez_compressed(temporary, **arrays)
|
||
|
|
with temporary.open("rb") as stream:
|
||
|
|
os.fsync(stream.fileno())
|
||
|
|
os.replace(temporary, path)
|
||
|
|
except Exception:
|
||
|
|
temporary.unlink(missing_ok=True)
|
||
|
|
raise
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_sample_arrays(
|
||
|
|
samples: Mapping[str, Any],
|
||
|
|
) -> tuple[dict[str, np.ndarray], int]:
|
||
|
|
if not isinstance(samples, Mapping) or not samples:
|
||
|
|
raise ContractError("trial samples must be a non-empty mapping")
|
||
|
|
arrays: dict[str, np.ndarray] = {}
|
||
|
|
sample_count: int | None = None
|
||
|
|
for name, value in samples.items():
|
||
|
|
if not isinstance(name, str) or not name:
|
||
|
|
raise ContractError("sample field names must be non-empty strings")
|
||
|
|
array = np.asarray(value)
|
||
|
|
if array.dtype == object:
|
||
|
|
raise ContractError(f"sample field {name!r} cannot have object dtype")
|
||
|
|
if array.ndim == 0:
|
||
|
|
raise ContractError(f"sample field {name!r} must have a sample axis")
|
||
|
|
if sample_count is None:
|
||
|
|
sample_count = int(array.shape[0])
|
||
|
|
elif array.shape[0] != sample_count:
|
||
|
|
raise ContractError(
|
||
|
|
f"sample field {name!r} has {array.shape[0]} rows, "
|
||
|
|
f"expected {sample_count}"
|
||
|
|
)
|
||
|
|
arrays[name] = array
|
||
|
|
assert sample_count is not None
|
||
|
|
if sample_count <= 0:
|
||
|
|
raise ContractError("trial samples must contain at least one row")
|
||
|
|
return arrays, sample_count
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class TrialPayload:
|
||
|
|
samples: Mapping[str, Any]
|
||
|
|
events: Sequence[Mapping[str, Any]] = field(default_factory=tuple)
|
||
|
|
metadata: Mapping[str, Any] = field(default_factory=dict)
|
||
|
|
|
||
|
|
|
||
|
|
class AtomicTrialWriter:
|
||
|
|
"""Commit one successful trial directory using a final atomic rename."""
|
||
|
|
|
||
|
|
def __init__(self, batch_dir: Path):
|
||
|
|
self.batch_dir = Path(batch_dir)
|
||
|
|
self.raw_dir = self.batch_dir / "raw"
|
||
|
|
self.raw_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
def trial_dir(self, trial_id: str) -> Path:
|
||
|
|
return self.raw_dir / trial_id
|
||
|
|
|
||
|
|
def write(self, trial: Mapping[str, Any], payload: TrialPayload) -> Path:
|
||
|
|
trial_id = str(trial["trial_id"])
|
||
|
|
destination = self.trial_dir(trial_id)
|
||
|
|
if destination.exists():
|
||
|
|
raise FileExistsError(f"trial artifact already exists: {destination}")
|
||
|
|
arrays, sample_count = normalize_sample_arrays(payload.samples)
|
||
|
|
temporary = self.raw_dir / f".{trial_id}.tmp.{uuid.uuid4().hex}"
|
||
|
|
temporary.mkdir(parents=False, exist_ok=False)
|
||
|
|
try:
|
||
|
|
sample_path = temporary / "samples.npz"
|
||
|
|
event_path = temporary / "events.jsonl"
|
||
|
|
atomic_write_npz(sample_path, arrays)
|
||
|
|
atomic_write_jsonl(event_path, list(payload.events))
|
||
|
|
fields = {
|
||
|
|
name: {"dtype": str(array.dtype), "shape": list(array.shape)}
|
||
|
|
for name, array in sorted(arrays.items())
|
||
|
|
}
|
||
|
|
manifest = {
|
||
|
|
"kind": "trial_manifest",
|
||
|
|
"schema_version": SCHEMA_VERSION,
|
||
|
|
"trial_id": trial_id,
|
||
|
|
"pair_id": trial["pair_id"],
|
||
|
|
"trial_spec_hash": trial["trial_spec_hash"],
|
||
|
|
"status": "completed",
|
||
|
|
"completed_utc": utc_now(),
|
||
|
|
"sample_count": sample_count,
|
||
|
|
"sample_fields": fields,
|
||
|
|
"storage_formats": {
|
||
|
|
"samples": STORAGE_FORMATS["samples"],
|
||
|
|
"events": STORAGE_FORMATS["events"],
|
||
|
|
},
|
||
|
|
"files": {
|
||
|
|
"samples.npz": file_sha256(sample_path),
|
||
|
|
"events.jsonl": file_sha256(event_path),
|
||
|
|
},
|
||
|
|
"metadata": to_jsonable(payload.metadata),
|
||
|
|
}
|
||
|
|
atomic_write_json(temporary / "trial_manifest.json", manifest)
|
||
|
|
os.replace(temporary, destination)
|
||
|
|
except Exception:
|
||
|
|
shutil.rmtree(temporary, ignore_errors=True)
|
||
|
|
raise
|
||
|
|
return destination
|
||
|
|
|
||
|
|
|
||
|
|
def record_trial_failure(
|
||
|
|
batch_dir: Path,
|
||
|
|
trial: Mapping[str, Any],
|
||
|
|
error: BaseException,
|
||
|
|
) -> Path:
|
||
|
|
failure_dir = Path(batch_dir) / "failures" / str(trial["trial_id"])
|
||
|
|
failure_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
path = failure_dir / f"attempt-{utc_now().replace(':', '')}-{uuid.uuid4().hex}.json"
|
||
|
|
atomic_write_json(
|
||
|
|
path,
|
||
|
|
{
|
||
|
|
"kind": "trial_failure",
|
||
|
|
"schema_version": SCHEMA_VERSION,
|
||
|
|
"trial_id": trial["trial_id"],
|
||
|
|
"pair_id": trial["pair_id"],
|
||
|
|
"trial_spec_hash": trial["trial_spec_hash"],
|
||
|
|
"recorded_utc": utc_now(),
|
||
|
|
"error_type": type(error).__name__,
|
||
|
|
"error_message": str(error),
|
||
|
|
},
|
||
|
|
)
|
||
|
|
return path
|