from __future__ import annotations import fcntl import json from dataclasses import asdict, dataclass, field, replace from pathlib import Path from typing import Any, Sequence import numpy as np from .config import DotCacheConfig from .decode_reference import decode_page from .encode import encode_page from .planner import PageModeSpec, observe_page, parse_page_mode_token from .types import Kind @dataclass(slots=True) class PageTraceRecord: source: str kind: Kind layer_id: int kv_head_id: int token_start: int token_age: int values: np.ndarray query: np.ndarray | None = None notes: list[str] = field(default_factory=list) def __post_init__(self) -> None: values = np.asarray(self.values, dtype=np.float32) if values.ndim != 2: raise ValueError("values must have shape [token_count, head_dim]") self.values = values if self.query is not None: query = np.asarray(self.query, dtype=np.float32) if query.ndim != 1: raise ValueError("query must have shape [head_dim]") if int(query.shape[0]) != int(values.shape[1]): raise ValueError("query head_dim must match values head_dim") self.query = query self.layer_id = int(self.layer_id) self.kv_head_id = int(self.kv_head_id) self.token_start = int(self.token_start) self.token_age = int(self.token_age) if self.token_age < 0: raise ValueError("token_age must be non-negative") @property def token_count(self) -> int: return int(self.values.shape[0]) @property def head_dim(self) -> int: return int(self.values.shape[1]) @property def stats(self): # pragma: no cover - thin wrapper return observe_page(self.values) def to_dict(self) -> dict[str, Any]: payload = { "source": self.source, "kind": self.kind, "layer_id": self.layer_id, "kv_head_id": self.kv_head_id, "token_start": self.token_start, "token_age": self.token_age, "token_count": self.token_count, "head_dim": self.head_dim, "notes": list(self.notes), "stats": asdict(self.stats), } if self.query is not None: payload["query_present"] = True return payload @dataclass(frozen=True, slots=True) class OracleThresholds: max_mean_abs_error_ratio: float = 0.10 max_max_abs_error_ratio: float = 1.00 max_token_p95_error_ratio: float = 0.25 max_score_max_abs_error: float | None = None min_score_topk_agreement: float | None = None @dataclass(slots=True) class OracleCandidateResult: candidate: str mode: str bits: int quant_scheme: str payload_bytes: int metadata_bytes: int total_bytes: int mean_abs_error: float max_abs_error: float token_p95_error: float score_max_abs_error: float | None score_topk_agreement: float | None safe: bool failure_reasons: list[str] = field(default_factory=list) def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass(slots=True) class OracleReplayResult: trace: dict[str, Any] thresholds: dict[str, Any] candidates: list[OracleCandidateResult] cheapest_safe_candidate: str | None def to_dict(self) -> dict[str, Any]: return { "trace": dict(self.trace), "thresholds": dict(self.thresholds), "cheapest_safe_candidate": self.cheapest_safe_candidate, "candidates": [candidate.to_dict() for candidate in self.candidates], } @dataclass(slots=True) class OracleBatchTraceResult: trace_path: str stage: str trace: dict[str, Any] cheapest_safe_candidate: str | None safe_candidate_count: int best_safe_total_bytes: int | None candidates: list[OracleCandidateResult] def to_dict(self) -> dict[str, Any]: return { "trace_path": self.trace_path, "stage": self.stage, "trace": dict(self.trace), "cheapest_safe_candidate": self.cheapest_safe_candidate, "safe_candidate_count": self.safe_candidate_count, "best_safe_total_bytes": self.best_safe_total_bytes, "candidates": [candidate.to_dict() for candidate in self.candidates], } @dataclass(slots=True) class OracleBatchReplayResult: manifest: dict[str, Any] sampling: dict[str, Any] thresholds: dict[str, Any] group_size: int tokens_per_page: int | None selected_trace_count: int selected_trace_paths: list[str] selected_trace_counts_by_stage: dict[str, int] selected_trace_counts_by_kind: dict[str, int] cheapest_safe_candidate_histogram: dict[str, int] failure_reason_histogram: dict[str, int] candidate_stats: dict[str, dict[str, Any]] summary_table_markdown: str traces: list[OracleBatchTraceResult] def to_dict(self) -> dict[str, Any]: return { "manifest": dict(self.manifest), "sampling": dict(self.sampling), "thresholds": dict(self.thresholds), "group_size": self.group_size, "tokens_per_page": self.tokens_per_page, "selected_trace_count": self.selected_trace_count, "selected_trace_paths": list(self.selected_trace_paths), "selected_trace_counts_by_stage": dict(self.selected_trace_counts_by_stage), "selected_trace_counts_by_kind": dict(self.selected_trace_counts_by_kind), "cheapest_safe_candidate_histogram": dict(self.cheapest_safe_candidate_histogram), "failure_reason_histogram": dict(self.failure_reason_histogram), "candidate_stats": dict(self.candidate_stats), "summary_table_markdown": self.summary_table_markdown, "traces": [trace.to_dict() for trace in self.traces], } @dataclass(slots=True) class OracleLabelRecord: trace_path: str stage: str prompt_family: str | None prompt_variant: str | None source: str kind: str layer_id: int kv_head_id: int token_start: int token_age: int token_count: int head_dim: int query_present: bool cheapest_safe_candidate: str | None safe_candidates: list[str] best_safe_total_bytes: int | None candidate_labels: list[dict[str, Any]] trace_stats: dict[str, Any] notes: list[str] def to_dict(self) -> dict[str, Any]: return { "trace_path": self.trace_path, "stage": self.stage, "prompt_family": self.prompt_family, "prompt_variant": self.prompt_variant, "source": self.source, "kind": self.kind, "layer_id": self.layer_id, "kv_head_id": self.kv_head_id, "token_start": self.token_start, "token_age": self.token_age, "token_count": self.token_count, "head_dim": self.head_dim, "query_present": self.query_present, "cheapest_safe_candidate": self.cheapest_safe_candidate, "safe_candidates": list(self.safe_candidates), "best_safe_total_bytes": self.best_safe_total_bytes, "candidate_labels": list(self.candidate_labels), "trace_stats": dict(self.trace_stats), "notes": list(self.notes), } @dataclass(slots=True) class OracleLabelingResult: labels: list[OracleLabelRecord] summary: dict[str, Any] def to_dict(self) -> dict[str, Any]: return { "summary": dict(self.summary), "labels": [label.to_dict() for label in self.labels], } @dataclass(slots=True) class OracleSelectorTrainingRow: trace_path: str source: str stage: str prompt_family: str | None prompt_variant: str | None kind: str layer_id: int layer_fraction: float kv_head_id: int kv_head_fraction: float token_start: int token_age: int token_count: int head_dim: int query_present: bool safe_candidate_count: int best_safe_total_bytes: int | None target_candidate: str | None target_present: bool trace_rms: float trace_abs_max: float trace_channel_range_mean: float trace_outlier_fraction: float age_per_token: float def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass(slots=True) class OracleSelectorCandidateTrainingRow: trace_path: str source: str stage: str prompt_family: str | None prompt_variant: str | None kind: str layer_id: int layer_fraction: float kv_head_id: int kv_head_fraction: float token_start: int token_age: int token_count: int head_dim: int query_present: bool safe_candidate_count: int best_safe_total_bytes: int | None target_candidate: str | None target_present: bool trace_rms: float trace_abs_max: float trace_channel_range_mean: float trace_outlier_fraction: float age_per_token: float candidate: str candidate_mode: str candidate_bits: int candidate_quant_scheme: str candidate_total_bytes: int candidate_payload_bytes: int candidate_metadata_bytes: int candidate_has_escape_dtype: bool candidate_safe: bool candidate_is_target: bool candidate_bytes_over_best_safe: int | None def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass(slots=True) class OracleDatasetSplitSummary: split_name: str holdout_prompt_families: list[str] holdout_prompt_variants: list[str] holdout_layers: list[int] train_trace_paths: list[str] test_trace_paths: list[str] train_label_count: int test_label_count: int train_selector_row_count: int test_selector_row_count: int train_selector_candidate_row_count: int test_selector_candidate_row_count: int train_prompt_family_histogram: dict[str, int] test_prompt_family_histogram: dict[str, int] train_prompt_variant_histogram: dict[str, int] test_prompt_variant_histogram: dict[str, int] train_layer_histogram: dict[str, int] test_layer_histogram: dict[str, int] def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass(slots=True) class OracleDatasetSplitManifestEntry: split_name: str split_dir: str split_summary_path: str holdout_prompt_families: list[str] holdout_prompt_variants: list[str] holdout_layers: list[int] train_label_count: int test_label_count: int annotations: dict[str, Any] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass(slots=True) class OracleDatasetSplitSuiteSpec: split_name: str output_subdir: str | None = None holdout_prompt_families: list[str] = field(default_factory=list) holdout_prompt_variants: list[str] = field(default_factory=list) holdout_layers: list[int] = field(default_factory=list) annotations: dict[str, Any] = field(default_factory=dict) def to_dict(self) -> dict[str, Any]: return asdict(self) @dataclass(slots=True) class OracleDatasetSplitSuiteResult: suite_name: str output_root: str manifest_path: str | None split_count: int split_names: list[str] split_dirs: list[str] splits: list[OracleDatasetSplitSummary] def to_dict(self) -> dict[str, Any]: return { "suite_name": self.suite_name, "output_root": self.output_root, "manifest_path": self.manifest_path, "split_count": self.split_count, "split_names": list(self.split_names), "split_dirs": list(self.split_dirs), "splits": [split.to_dict() for split in self.splits], } def default_candidate_specs(kind: Kind) -> tuple[PageModeSpec, ...]: if kind == "K": tokens = ( "M0/affine/2", "M0/affine/3", "M0/affine/4", "M2/sketch/4", "M4/project/4", "M3/affine/4/float16", "T3/turbo3/3", ) else: tokens = ( "M0/affine/2", "M0/affine/3", "M0/affine/4", "M1/lut/4", "M3/affine/4/float16", "T3/turbo3/3", ) return tuple(parse_page_mode_token(token) for token in tokens) def save_page_trace(record: PageTraceRecord, path: str | Path) -> None: target = Path(path) target.parent.mkdir(parents=True, exist_ok=True) arrays: dict[str, Any] = { "values": np.asarray(record.values, dtype=np.float32), "metadata_json": np.array(json.dumps(record.to_dict()), dtype=np.str_), } if record.query is not None: arrays["query"] = np.asarray(record.query, dtype=np.float32) np.savez_compressed(target, **arrays) def load_page_trace(path: str | Path) -> PageTraceRecord: with np.load(Path(path), allow_pickle=False) as payload: metadata = json.loads(str(payload["metadata_json"])) query = payload["query"] if "query" in payload else None return PageTraceRecord( source=str(metadata.get("source", "")), kind=str(metadata["kind"]), layer_id=int(metadata["layer_id"]), kv_head_id=int(metadata["kv_head_id"]), token_start=int(metadata["token_start"]), token_age=int(metadata["token_age"]), values=np.asarray(payload["values"], dtype=np.float32), query=None if query is None else np.asarray(query, dtype=np.float32), notes=list(metadata.get("notes", [])), ) def _config_for_candidate( base_config: DotCacheConfig, *, kind: Kind, candidate: PageModeSpec, ) -> DotCacheConfig: if kind == "K": return replace( base_config, bits_k=int(candidate.bits), default_mode_k=str(candidate.mode), quant_scheme_k=str(candidate.quant_scheme), escape_dtype=str(candidate.escape_dtype or base_config.escape_dtype), ) return replace( base_config, bits_v=int(candidate.bits), default_mode_v=str(candidate.mode), quant_scheme_v=str(candidate.quant_scheme), escape_dtype=str(candidate.escape_dtype or base_config.escape_dtype), ) def _score_metrics( values: np.ndarray, reconstructed: np.ndarray, query: np.ndarray | None, ) -> tuple[float | None, float | None]: if query is None: return None, None dense_scores = values @ query approx_scores = reconstructed @ query score_max_abs_error = float(np.max(np.abs(dense_scores - approx_scores))) top_k = min(4, int(values.shape[0])) dense_top = set(np.argsort(dense_scores)[-top_k:].tolist()) approx_top = set(np.argsort(approx_scores)[-top_k:].tolist()) topk_agreement = float(len(dense_top & approx_top) / max(top_k, 1)) return score_max_abs_error, topk_agreement def evaluate_page_candidate( record: PageTraceRecord, *, base_config: DotCacheConfig, candidate: PageModeSpec, thresholds: OracleThresholds, ) -> OracleCandidateResult: candidate_config = _config_for_candidate(base_config, kind=record.kind, candidate=candidate) encoded = encode_page( record.values, candidate_config, kind=record.kind, layer_id=record.layer_id, kv_head_id=record.kv_head_id, token_start=record.token_start, page_mode=candidate, ) reconstructed = decode_page(encoded) abs_error = np.abs(record.values - reconstructed) token_max_error = np.max(abs_error, axis=1) page_rms = max(float(record.stats.rms), 1e-6) mean_abs_error = float(np.mean(abs_error)) max_abs_error = float(np.max(abs_error)) token_p95_error = float(np.percentile(token_max_error, 95)) score_max_abs_error, score_topk_agreement = _score_metrics(record.values, reconstructed, record.query) failure_reasons: list[str] = [] if mean_abs_error > thresholds.max_mean_abs_error_ratio * page_rms: failure_reasons.append("mean_abs_error") if max_abs_error > thresholds.max_max_abs_error_ratio * page_rms: failure_reasons.append("max_abs_error") if token_p95_error > thresholds.max_token_p95_error_ratio * page_rms: failure_reasons.append("token_p95_error") if thresholds.max_score_max_abs_error is not None and score_max_abs_error is not None: if score_max_abs_error > float(thresholds.max_score_max_abs_error): failure_reasons.append("score_max_abs_error") if thresholds.min_score_topk_agreement is not None and score_topk_agreement is not None: if score_topk_agreement < float(thresholds.min_score_topk_agreement): failure_reasons.append("score_topk_agreement") return OracleCandidateResult( candidate=_format_candidate_token(candidate), mode=str(candidate.mode), bits=int(candidate.bits), quant_scheme=str(candidate.quant_scheme), payload_bytes=int(encoded.payload_nbytes), metadata_bytes=int(encoded.metadata_nbytes), total_bytes=int(encoded.total_nbytes), mean_abs_error=mean_abs_error, max_abs_error=max_abs_error, token_p95_error=token_p95_error, score_max_abs_error=score_max_abs_error, score_topk_agreement=score_topk_agreement, safe=not failure_reasons, failure_reasons=failure_reasons, ) def run_oracle_replay( record: PageTraceRecord, *, base_config: DotCacheConfig, candidates: Sequence[PageModeSpec] | None = None, thresholds: OracleThresholds | None = None, ) -> OracleReplayResult: resolved_thresholds = thresholds or OracleThresholds() resolved_candidates = tuple(candidates or default_candidate_specs(record.kind)) candidate_results = [ evaluate_page_candidate( record, base_config=base_config, candidate=candidate, thresholds=resolved_thresholds, ) for candidate in resolved_candidates ] safe_candidates = [candidate for candidate in candidate_results if candidate.safe] safe_candidates.sort(key=lambda item: (item.total_bytes, item.mean_abs_error, item.max_abs_error)) return OracleReplayResult( trace=record.to_dict(), thresholds=asdict(resolved_thresholds), candidates=candidate_results, cheapest_safe_candidate=safe_candidates[0].candidate if safe_candidates else None, ) def load_page_trace_manifest(path: str | Path) -> dict[str, Any]: manifest_path = Path(path) return json.loads(manifest_path.read_text(encoding="utf-8")) def save_page_trace_manifest(manifest: dict[str, Any], path: str | Path) -> None: target = Path(path) target.parent.mkdir(parents=True, exist_ok=True) target.write_text(json.dumps(manifest, sort_keys=True, indent=2) + "\n", encoding="utf-8") def merge_page_trace_manifests( manifests: Sequence[dict[str, Any] | str | Path], *, output_dir: str | Path | None = None, source: str = "merged_page_trace_manifest", ) -> dict[str, Any]: resolved_manifests = [ load_page_trace_manifest(manifest) if isinstance(manifest, (str, Path)) else dict(manifest) for manifest in manifests ] page_trace_paths: list[str] = [] counts_by_kind: dict[str, int] = {} counts_by_stage: dict[str, int] = {} counts_by_layer: dict[str, int] = {} tokens_per_page_values: set[int] = set() kinds_union: set[str] = set() member_output_dirs: list[str] = [] member_sources: list[str] = [] for manifest in resolved_manifests: page_trace_paths.extend(str(path) for path in manifest.get("page_trace_paths", [])) for key, value in dict(manifest.get("page_trace_counts_by_kind", {})).items(): counts_by_kind[str(key)] = counts_by_kind.get(str(key), 0) + int(value) for key, value in dict(manifest.get("page_trace_counts_by_stage", {})).items(): counts_by_stage[str(key)] = counts_by_stage.get(str(key), 0) + int(value) for key, value in dict(manifest.get("page_trace_counts_by_layer", {})).items(): counts_by_layer[str(key)] = counts_by_layer.get(str(key), 0) + int(value) if "tokens_per_page" in manifest: tokens_per_page_values.add(int(manifest["tokens_per_page"])) for kind in manifest.get("kinds", []): kinds_union.add(str(kind)) if "output_dir" in manifest: member_output_dirs.append(str(manifest["output_dir"])) if "source" in manifest: member_sources.append(str(manifest["source"])) merged_manifest = { "output_dir": None if output_dir is None else str(output_dir), "page_trace_count": len(page_trace_paths), "page_trace_paths": page_trace_paths, "page_trace_counts_by_kind": dict(sorted(counts_by_kind.items())), "page_trace_counts_by_stage": dict(sorted(counts_by_stage.items())), "page_trace_counts_by_layer": dict(sorted(counts_by_layer.items())), "tokens_per_page": None if len(tokens_per_page_values) != 1 else next(iter(tokens_per_page_values)), "tokens_per_page_values": sorted(tokens_per_page_values), "kinds": sorted(kinds_union), "source": source, "member_output_dirs": member_output_dirs, "member_sources": sorted(set(member_sources)), "member_manifest_count": len(resolved_manifests), } return merged_manifest def select_page_trace_paths( manifest: dict[str, Any] | str | Path, *, max_traces: int | None = None, max_per_stage_kind: int | None = None, seed: int = 0, kinds: Sequence[str] | None = None, stages: Sequence[str] | None = None, layer_ids: Sequence[int] | None = None, ) -> list[str]: manifest_payload = load_page_trace_manifest(manifest) if isinstance(manifest, (str, Path)) else dict(manifest) trace_paths = [str(path) for path in manifest_payload.get("page_trace_paths", [])] normalized_kinds = None if kinds is None else {str(kind).upper() for kind in kinds} normalized_stages = None if stages is None else {str(stage).lower() for stage in stages} normalized_layers = None if layer_ids is None else {int(layer_id) for layer_id in layer_ids} rng = np.random.default_rng(int(seed)) grouped_matches: dict[tuple[str, str], list[str]] = {} for trace_path in trace_paths: record = load_page_trace(trace_path) stage = _trace_stage(record) kind = str(record.kind).upper() if normalized_kinds is not None and kind not in normalized_kinds: continue if normalized_stages is not None and stage not in normalized_stages: continue if normalized_layers is not None and int(record.layer_id) not in normalized_layers: continue grouped_matches.setdefault((stage, kind), []).append(str(trace_path)) selected_paths: list[str] = [] for group_key in sorted(grouped_matches): group_paths = list(grouped_matches[group_key]) if len(group_paths) > 1: order = rng.permutation(len(group_paths)).tolist() group_paths = [group_paths[index] for index in order] if max_per_stage_kind is not None: group_paths = group_paths[: max(int(max_per_stage_kind), 0)] selected_paths.extend(group_paths) if len(selected_paths) > 1: order = rng.permutation(len(selected_paths)).tolist() selected_paths = [selected_paths[index] for index in order] if max_traces is not None: selected_paths = selected_paths[: max(int(max_traces), 0)] return selected_paths def run_oracle_batch_replay( manifest: dict[str, Any] | str | Path, *, group_size: int, tokens_per_page: int | None = None, candidates: Sequence[PageModeSpec] | None = None, thresholds: OracleThresholds | None = None, max_traces: int | None = None, max_per_stage_kind: int | None = None, seed: int = 0, kinds: Sequence[str] | None = None, stages: Sequence[str] | None = None, layer_ids: Sequence[int] | None = None, ) -> OracleBatchReplayResult: manifest_payload = load_page_trace_manifest(manifest) if isinstance(manifest, (str, Path)) else dict(manifest) resolved_thresholds = thresholds or OracleThresholds() selected_paths = select_page_trace_paths( manifest_payload, max_traces=max_traces, max_per_stage_kind=max_per_stage_kind, seed=seed, kinds=kinds, stages=stages, layer_ids=layer_ids, ) traces: list[OracleBatchTraceResult] = [] counts_by_stage: dict[str, int] = {} counts_by_kind: dict[str, int] = {} selected_histogram: dict[str, int] = {} failure_reason_histogram: dict[str, int] = {} candidate_stats_raw: dict[str, dict[str, float]] = {} for trace_path in selected_paths: record = load_page_trace(trace_path) stage = _trace_stage(record) counts_by_stage[stage] = counts_by_stage.get(stage, 0) + 1 counts_by_kind[record.kind] = counts_by_kind.get(record.kind, 0) + 1 replay = run_oracle_replay( record, base_config=DotCacheConfig( head_dim=record.head_dim, group_size=int(group_size), tokens_per_page=int(tokens_per_page or record.token_count), ), candidates=candidates, thresholds=resolved_thresholds, ) safe_candidates = [candidate for candidate in replay.candidates if candidate.safe] if replay.cheapest_safe_candidate is not None: selected_histogram[replay.cheapest_safe_candidate] = selected_histogram.get(replay.cheapest_safe_candidate, 0) + 1 for candidate in replay.candidates: stats = candidate_stats_raw.setdefault( candidate.candidate, { "count": 0.0, "safe_count": 0.0, "sum_total_bytes": 0.0, "sum_mean_abs_error": 0.0, "sum_token_p95_error": 0.0, }, ) stats["count"] += 1.0 stats["safe_count"] += 1.0 if candidate.safe else 0.0 stats["sum_total_bytes"] += float(candidate.total_bytes) stats["sum_mean_abs_error"] += float(candidate.mean_abs_error) stats["sum_token_p95_error"] += float(candidate.token_p95_error) for reason in candidate.failure_reasons: failure_reason_histogram[reason] = failure_reason_histogram.get(reason, 0) + 1 traces.append( OracleBatchTraceResult( trace_path=str(trace_path), stage=stage, trace=replay.trace, cheapest_safe_candidate=replay.cheapest_safe_candidate, safe_candidate_count=len(safe_candidates), best_safe_total_bytes=min((candidate.total_bytes for candidate in safe_candidates), default=None), candidates=replay.candidates, ) ) candidate_stats: dict[str, dict[str, Any]] = {} for candidate, stats in sorted(candidate_stats_raw.items()): count = max(int(stats["count"]), 1) candidate_stats[candidate] = { "count": int(stats["count"]), "safe_count": int(stats["safe_count"]), "safe_rate": float(stats["safe_count"] / count), "mean_total_bytes": float(stats["sum_total_bytes"] / count), "mean_abs_error": float(stats["sum_mean_abs_error"] / count), "mean_token_p95_error": float(stats["sum_token_p95_error"] / count), "selected_count": int(selected_histogram.get(candidate, 0)), } return OracleBatchReplayResult( manifest=manifest_payload, sampling={ "seed": int(seed), "max_traces": None if max_traces is None else int(max_traces), "max_per_stage_kind": None if max_per_stage_kind is None else int(max_per_stage_kind), "kinds": None if kinds is None else [str(kind).upper() for kind in kinds], "stages": None if stages is None else [str(stage).lower() for stage in stages], "layer_ids": None if layer_ids is None else [int(layer_id) for layer_id in layer_ids], }, thresholds=asdict(resolved_thresholds), group_size=int(group_size), tokens_per_page=None if tokens_per_page is None else int(tokens_per_page), selected_trace_count=len(traces), selected_trace_paths=list(selected_paths), selected_trace_counts_by_stage=dict(sorted(counts_by_stage.items())), selected_trace_counts_by_kind=dict(sorted(counts_by_kind.items())), cheapest_safe_candidate_histogram=dict(sorted(selected_histogram.items())), failure_reason_histogram=dict(sorted(failure_reason_histogram.items())), candidate_stats=candidate_stats, summary_table_markdown=_render_batch_summary_table(candidate_stats), traces=traces, ) def build_oracle_label_records(batch_result: OracleBatchReplayResult) -> list[OracleLabelRecord]: labels: list[OracleLabelRecord] = [] for trace_result in batch_result.traces: safe_candidates = sorted( (candidate for candidate in trace_result.candidates if candidate.safe), key=lambda item: (item.total_bytes, item.mean_abs_error, item.max_abs_error, item.candidate), ) trace = dict(trace_result.trace) labels.append( OracleLabelRecord( trace_path=str(trace_result.trace_path), stage=str(trace_result.stage), prompt_family=_infer_prompt_family_from_trace_path(trace_result.trace_path), prompt_variant=_infer_prompt_variant_from_trace_path(trace_result.trace_path), source=str(trace.get("source", "")), kind=str(trace["kind"]), layer_id=int(trace["layer_id"]), kv_head_id=int(trace["kv_head_id"]), token_start=int(trace["token_start"]), token_age=int(trace["token_age"]), token_count=int(trace["token_count"]), head_dim=int(trace["head_dim"]), query_present=bool(trace.get("query_present", False)), cheapest_safe_candidate=trace_result.cheapest_safe_candidate, safe_candidates=[candidate.candidate for candidate in safe_candidates], best_safe_total_bytes=trace_result.best_safe_total_bytes, candidate_labels=[candidate.to_dict() for candidate in trace_result.candidates], trace_stats=dict(trace.get("stats", {})), notes=list(trace.get("notes", [])), ) ) return labels def run_oracle_labeling( manifest: dict[str, Any] | str | Path, *, group_size: int, tokens_per_page: int | None = None, candidates: Sequence[PageModeSpec] | None = None, thresholds: OracleThresholds | None = None, max_traces: int | None = None, max_per_stage_kind: int | None = None, seed: int = 0, kinds: Sequence[str] | None = None, stages: Sequence[str] | None = None, layer_ids: Sequence[int] | None = None, ) -> OracleLabelingResult: batch_result = run_oracle_batch_replay( manifest, group_size=group_size, tokens_per_page=tokens_per_page, candidates=candidates, thresholds=thresholds, max_traces=max_traces, max_per_stage_kind=max_per_stage_kind, seed=seed, kinds=kinds, stages=stages, layer_ids=layer_ids, ) labels = build_oracle_label_records(batch_result) summary = { "manifest": dict(batch_result.manifest), "sampling": dict(batch_result.sampling), "thresholds": dict(batch_result.thresholds), "group_size": int(batch_result.group_size), "tokens_per_page": batch_result.tokens_per_page, "label_count": len(labels), "selected_trace_counts_by_stage": dict(batch_result.selected_trace_counts_by_stage), "selected_trace_counts_by_kind": dict(batch_result.selected_trace_counts_by_kind), "cheapest_safe_candidate_histogram": dict(batch_result.cheapest_safe_candidate_histogram), "failure_reason_histogram": dict(batch_result.failure_reason_histogram), "candidate_stats": dict(batch_result.candidate_stats), "summary_table_markdown": batch_result.summary_table_markdown, } return OracleLabelingResult(labels=labels, summary=summary) def save_oracle_labels( labeling_result: OracleLabelingResult, *, labels_path: str | Path, summary_path: str | Path | None = None, ) -> None: labels_target = Path(labels_path) labels_target.parent.mkdir(parents=True, exist_ok=True) with labels_target.open("w", encoding="utf-8") as handle: for label in labeling_result.labels: handle.write(json.dumps(label.to_dict(), sort_keys=True) + "\n") if summary_path is not None: summary_target = Path(summary_path) summary_target.parent.mkdir(parents=True, exist_ok=True) summary_target.write_text(json.dumps(labeling_result.summary, sort_keys=True, indent=2) + "\n", encoding="utf-8") def build_selector_training_rows(labels: Sequence[OracleLabelRecord]) -> list[OracleSelectorTrainingRow]: if not labels: return [] max_layer_id = max(label.layer_id for label in labels) max_kv_head_id = max(label.kv_head_id for label in labels) rows: list[OracleSelectorTrainingRow] = [] for label in labels: stats = dict(label.trace_stats) rows.append( OracleSelectorTrainingRow( trace_path=str(label.trace_path), source=str(label.source), stage=str(label.stage), prompt_family=label.prompt_family, prompt_variant=label.prompt_variant, kind=str(label.kind), layer_id=int(label.layer_id), layer_fraction=float(label.layer_id / max(max_layer_id, 1)), kv_head_id=int(label.kv_head_id), kv_head_fraction=float(label.kv_head_id / max(max_kv_head_id, 1)), token_start=int(label.token_start), token_age=int(label.token_age), token_count=int(label.token_count), head_dim=int(label.head_dim), query_present=bool(label.query_present), safe_candidate_count=len(label.safe_candidates), best_safe_total_bytes=None if label.best_safe_total_bytes is None else int(label.best_safe_total_bytes), target_candidate=label.cheapest_safe_candidate, target_present=label.cheapest_safe_candidate is not None, trace_rms=float(stats.get("rms", 0.0)), trace_abs_max=float(stats.get("abs_max", 0.0)), trace_channel_range_mean=float(stats.get("channel_range_mean", 0.0)), trace_outlier_fraction=float(stats.get("outlier_fraction", 0.0)), age_per_token=float(label.token_age / max(label.token_count, 1)), ) ) return rows def save_selector_training_rows( rows: Sequence[OracleSelectorTrainingRow], path: str | Path, ) -> None: target = Path(path) target.parent.mkdir(parents=True, exist_ok=True) with target.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row.to_dict(), sort_keys=True) + "\n") def build_selector_candidate_training_rows(labels: Sequence[OracleLabelRecord]) -> list[OracleSelectorCandidateTrainingRow]: if not labels: return [] max_layer_id = max(label.layer_id for label in labels) max_kv_head_id = max(label.kv_head_id for label in labels) rows: list[OracleSelectorCandidateTrainingRow] = [] for label in labels: stats = dict(label.trace_stats) for candidate_payload in label.candidate_labels: candidate = str(candidate_payload["candidate"]) rows.append( OracleSelectorCandidateTrainingRow( trace_path=str(label.trace_path), source=str(label.source), stage=str(label.stage), prompt_family=label.prompt_family, prompt_variant=label.prompt_variant, kind=str(label.kind), layer_id=int(label.layer_id), layer_fraction=float(label.layer_id / max(max_layer_id, 1)), kv_head_id=int(label.kv_head_id), kv_head_fraction=float(label.kv_head_id / max(max_kv_head_id, 1)), token_start=int(label.token_start), token_age=int(label.token_age), token_count=int(label.token_count), head_dim=int(label.head_dim), query_present=bool(label.query_present), safe_candidate_count=len(label.safe_candidates), best_safe_total_bytes=None if label.best_safe_total_bytes is None else int(label.best_safe_total_bytes), target_candidate=label.cheapest_safe_candidate, target_present=label.cheapest_safe_candidate is not None, trace_rms=float(stats.get("rms", 0.0)), trace_abs_max=float(stats.get("abs_max", 0.0)), trace_channel_range_mean=float(stats.get("channel_range_mean", 0.0)), trace_outlier_fraction=float(stats.get("outlier_fraction", 0.0)), age_per_token=float(label.token_age / max(label.token_count, 1)), candidate=candidate, candidate_mode=str(candidate_payload["mode"]), candidate_bits=int(candidate_payload["bits"]), candidate_quant_scheme=str(candidate_payload["quant_scheme"]), candidate_total_bytes=int(candidate_payload["total_bytes"]), candidate_payload_bytes=int(candidate_payload["payload_bytes"]), candidate_metadata_bytes=int(candidate_payload["metadata_bytes"]), candidate_has_escape_dtype=len(candidate.split("/")) >= 4, candidate_safe=bool(candidate_payload.get("safe", False)), candidate_is_target=bool(label.cheapest_safe_candidate is not None and candidate == label.cheapest_safe_candidate), candidate_bytes_over_best_safe=( None if label.best_safe_total_bytes is None else int(candidate_payload["total_bytes"]) - int(label.best_safe_total_bytes) ), ) ) return rows def save_selector_candidate_training_rows( rows: Sequence[OracleSelectorCandidateTrainingRow], path: str | Path, ) -> None: target = Path(path) target.parent.mkdir(parents=True, exist_ok=True) with target.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row.to_dict(), sort_keys=True) + "\n") def load_oracle_label_records(path: str | Path) -> list[OracleLabelRecord]: records: list[OracleLabelRecord] = [] for payload in _load_jsonl_records(path): records.append( OracleLabelRecord( trace_path=str(payload["trace_path"]), stage=str(payload["stage"]), prompt_family=None if payload.get("prompt_family") in (None, "") else str(payload.get("prompt_family")), prompt_variant=None if payload.get("prompt_variant") in (None, "") else str(payload.get("prompt_variant")), source=str(payload["source"]), kind=str(payload["kind"]), layer_id=int(payload["layer_id"]), kv_head_id=int(payload["kv_head_id"]), token_start=int(payload["token_start"]), token_age=int(payload["token_age"]), token_count=int(payload["token_count"]), head_dim=int(payload["head_dim"]), query_present=bool(payload["query_present"]), cheapest_safe_candidate=None if payload.get("cheapest_safe_candidate") in (None, "") else str(payload.get("cheapest_safe_candidate")), safe_candidates=[str(candidate) for candidate in payload.get("safe_candidates", [])], best_safe_total_bytes=None if payload.get("best_safe_total_bytes") is None else int(payload["best_safe_total_bytes"]), candidate_labels=[dict(candidate) for candidate in payload.get("candidate_labels", [])], trace_stats=dict(payload.get("trace_stats", {})), notes=[str(note) for note in payload.get("notes", [])], ) ) return records def materialize_oracle_dataset_split( *, labels_path: str | Path, selector_dataset_path: str | Path, output_dir: str | Path, selector_candidate_dataset_path: str | Path | None = None, holdout_prompt_families: Sequence[str] | None = None, holdout_prompt_variants: Sequence[str] | None = None, holdout_layers: Sequence[int] | None = None, split_name: str = "heldout_split", manifest_path: str | Path | None = None, annotations: dict[str, Any] | None = None, ) -> OracleDatasetSplitSummary: normalized_holdout_families = { token for token in (_normalize_split_token(value) for value in (holdout_prompt_families or ())) if token is not None } normalized_holdout_variants = { token for token in (_normalize_split_token(value) for value in (holdout_prompt_variants or ())) if token is not None } normalized_holdout_layers = {int(layer) for layer in (holdout_layers or ())} if not normalized_holdout_families and not normalized_holdout_variants and not normalized_holdout_layers: raise ValueError("at least one holdout filter must be provided") labels = load_oracle_label_records(labels_path) selector_rows = _load_jsonl_records(selector_dataset_path) selector_candidate_rows = ( [] if selector_candidate_dataset_path is None else _load_jsonl_records(selector_candidate_dataset_path) ) label_by_trace_path = {label.trace_path: label for label in labels} selector_trace_paths = {str(row["trace_path"]) for row in selector_rows} if selector_candidate_rows: selector_trace_paths |= {str(row["trace_path"]) for row in selector_candidate_rows} missing_trace_paths = sorted(selector_trace_paths - set(label_by_trace_path)) if missing_trace_paths: raise ValueError(f"dataset rows are missing matching labels for trace paths: {missing_trace_paths[:3]}") test_trace_paths = sorted( trace_path for trace_path, label in label_by_trace_path.items() if _matches_holdout_filters( prompt_family=label.prompt_family, prompt_variant=label.prompt_variant, layer_id=label.layer_id, holdout_prompt_families=normalized_holdout_families, holdout_prompt_variants=normalized_holdout_variants, holdout_layers=normalized_holdout_layers, ) ) train_trace_paths = sorted(trace_path for trace_path in label_by_trace_path if trace_path not in set(test_trace_paths)) if not train_trace_paths or not test_trace_paths: raise ValueError("split would produce an empty train or test partition") test_trace_path_set = set(test_trace_paths) train_trace_path_set = set(train_trace_paths) train_labels = [label for label in labels if label.trace_path in train_trace_path_set] test_labels = [label for label in labels if label.trace_path in test_trace_path_set] train_selector_rows = [row for row in selector_rows if str(row["trace_path"]) in train_trace_path_set] test_selector_rows = [row for row in selector_rows if str(row["trace_path"]) in test_trace_path_set] train_selector_candidate_rows = [row for row in selector_candidate_rows if str(row["trace_path"]) in train_trace_path_set] test_selector_candidate_rows = [row for row in selector_candidate_rows if str(row["trace_path"]) in test_trace_path_set] output_root = Path(output_dir) train_dir = output_root / "train" test_dir = output_root / "test" train_dir.mkdir(parents=True, exist_ok=True) test_dir.mkdir(parents=True, exist_ok=True) _save_jsonl_records(train_dir / "labels.jsonl", [label.to_dict() for label in train_labels]) _save_jsonl_records(test_dir / "labels.jsonl", [label.to_dict() for label in test_labels]) _save_jsonl_records(train_dir / "selector_dataset.jsonl", train_selector_rows) _save_jsonl_records(test_dir / "selector_dataset.jsonl", test_selector_rows) _save_schema(train_dir / "selector_schema.json", train_selector_rows) _save_schema(test_dir / "selector_schema.json", test_selector_rows) if selector_candidate_dataset_path is not None: _save_jsonl_records(train_dir / "selector_candidate_dataset.jsonl", train_selector_candidate_rows) _save_jsonl_records(test_dir / "selector_candidate_dataset.jsonl", test_selector_candidate_rows) _save_schema(train_dir / "selector_candidate_schema.json", train_selector_candidate_rows) _save_schema(test_dir / "selector_candidate_schema.json", test_selector_candidate_rows) summary = OracleDatasetSplitSummary( split_name=str(split_name), holdout_prompt_families=sorted(normalized_holdout_families), holdout_prompt_variants=sorted(normalized_holdout_variants), holdout_layers=sorted(normalized_holdout_layers), train_trace_paths=train_trace_paths, test_trace_paths=test_trace_paths, train_label_count=len(train_labels), test_label_count=len(test_labels), train_selector_row_count=len(train_selector_rows), test_selector_row_count=len(test_selector_rows), train_selector_candidate_row_count=len(train_selector_candidate_rows), test_selector_candidate_row_count=len(test_selector_candidate_rows), train_prompt_family_histogram=_label_histogram(train_labels, field_name="prompt_family"), test_prompt_family_histogram=_label_histogram(test_labels, field_name="prompt_family"), train_prompt_variant_histogram=_label_histogram(train_labels, field_name="prompt_variant"), test_prompt_variant_histogram=_label_histogram(test_labels, field_name="prompt_variant"), train_layer_histogram=_label_histogram(train_labels, field_name="layer_id"), test_layer_histogram=_label_histogram(test_labels, field_name="layer_id"), ) (output_root / "split_summary.json").write_text(json.dumps(summary.to_dict(), sort_keys=True, indent=2) + "\n", encoding="utf-8") if manifest_path is not None: upsert_oracle_dataset_split_manifest_entry( manifest_path, split_dir=output_root, summary=summary, annotations=annotations, ) return summary def materialize_oracle_dataset_split_suite( *, labels_path: str | Path, selector_dataset_path: str | Path, output_root: str | Path, suite_specs: Sequence[OracleDatasetSplitSuiteSpec | dict[str, Any]], selector_candidate_dataset_path: str | Path | None = None, suite_name: str = "selector_split_suite", manifest_path: str | Path | None = None, overwrite_manifest: bool = True, ) -> OracleDatasetSplitSuiteResult: resolved_specs = [_coerce_split_suite_spec(spec) for spec in suite_specs] output_root_path = Path(output_root) output_root_path.mkdir(parents=True, exist_ok=True) if manifest_path is not None and overwrite_manifest: save_oracle_dataset_split_manifest( { "manifest_version": 1, "suite_name": str(suite_name), "split_count": 0, "splits": [], }, manifest_path, ) split_summaries: list[OracleDatasetSplitSummary] = [] split_dirs: list[str] = [] for spec in resolved_specs: split_dir = output_root_path / (spec.output_subdir or spec.split_name) summary = materialize_oracle_dataset_split( labels_path=labels_path, selector_dataset_path=selector_dataset_path, selector_candidate_dataset_path=selector_candidate_dataset_path, output_dir=split_dir, holdout_prompt_families=tuple(spec.holdout_prompt_families) or None, holdout_prompt_variants=tuple(spec.holdout_prompt_variants) or None, holdout_layers=tuple(spec.holdout_layers) or None, split_name=spec.split_name, manifest_path=manifest_path, annotations=spec.annotations, ) split_summaries.append(summary) split_dirs.append(str(split_dir)) result = OracleDatasetSplitSuiteResult( suite_name=str(suite_name), output_root=str(output_root_path), manifest_path=None if manifest_path is None else str(manifest_path), split_count=len(split_summaries), split_names=[summary.split_name for summary in split_summaries], split_dirs=split_dirs, splits=split_summaries, ) (output_root_path / "split_suite_summary.json").write_text( json.dumps(result.to_dict(), sort_keys=True, indent=2) + "\n", encoding="utf-8", ) return result def _format_candidate_token(candidate: PageModeSpec) -> str: token = f"{candidate.mode}/{candidate.quant_scheme}/{candidate.bits}" if candidate.escape_dtype is not None: token = f"{token}/{candidate.escape_dtype}" return token def _trace_stage(record: PageTraceRecord) -> str: for note in record.notes: if note.startswith("stage="): return note.split("=", 1)[1].lower() return "unknown" def _infer_prompt_family_from_trace_path(trace_path: str | Path) -> str | None: for part in Path(trace_path).parts: if part.startswith("family-"): family = part.removeprefix("family-") if "_variant-" in family: family, _ = family.split("_variant-", 1) elif "_prompt" in family: family, _ = family.split("_prompt", 1) return family or None return None def _infer_prompt_variant_from_trace_path(trace_path: str | Path) -> str | None: for part in Path(trace_path).parts: if "_variant-" in part: _, rest = part.split("_variant-", 1) variant, *_ = rest.split("_prompt", 1) return variant return None def _load_jsonl_records(path: str | Path) -> list[dict[str, Any]]: records: list[dict[str, Any]] = [] with Path(path).open("r", encoding="utf-8") as handle: for line in handle: if not line.strip(): continue records.append(json.loads(line)) return records def _save_jsonl_records(path: str | Path, records: Sequence[dict[str, Any]]) -> None: target = Path(path) target.parent.mkdir(parents=True, exist_ok=True) with target.open("w", encoding="utf-8") as handle: for record in records: handle.write(json.dumps(record, sort_keys=True) + "\n") def _save_schema(path: str | Path, records: Sequence[dict[str, Any]]) -> None: schema = { "row_count": len(records), "fields": sorted(records[0].keys()) if records else [], } target = Path(path) target.parent.mkdir(parents=True, exist_ok=True) target.write_text(json.dumps(schema, sort_keys=True, indent=2) + "\n", encoding="utf-8") def _normalize_split_token(value: Any) -> str | None: if value in (None, ""): return None return str(value).strip().lower() or None def _coerce_split_suite_spec(spec: OracleDatasetSplitSuiteSpec | dict[str, Any]) -> OracleDatasetSplitSuiteSpec: if isinstance(spec, OracleDatasetSplitSuiteSpec): return spec payload = dict(spec) split_name = str(payload["split_name"]) output_subdir = payload.get("output_subdir") return OracleDatasetSplitSuiteSpec( split_name=split_name, output_subdir=None if output_subdir in (None, "") else str(output_subdir), holdout_prompt_families=[str(item) for item in payload.get("holdout_prompt_families", [])], holdout_prompt_variants=[str(item) for item in payload.get("holdout_prompt_variants", [])], holdout_layers=[int(item) for item in payload.get("holdout_layers", [])], annotations=dict(payload.get("annotations", {})), ) def _matches_holdout_filters( *, prompt_family: str | None, prompt_variant: str | None, layer_id: int, holdout_prompt_families: set[str], holdout_prompt_variants: set[str], holdout_layers: set[int], ) -> bool: if holdout_prompt_families and _normalize_split_token(prompt_family) not in holdout_prompt_families: return False if holdout_prompt_variants and _normalize_split_token(prompt_variant) not in holdout_prompt_variants: return False if holdout_layers and int(layer_id) not in holdout_layers: return False return True def _label_histogram(labels: Sequence[OracleLabelRecord], *, field_name: str) -> dict[str, int]: histogram: dict[str, int] = {} for label in labels: raw_value = getattr(label, field_name) key = "null" if raw_value is None else str(raw_value) histogram[key] = histogram.get(key, 0) + 1 return dict(sorted(histogram.items())) def load_oracle_dataset_split_manifest(path: str | Path) -> dict[str, Any]: manifest_path = Path(path) if not manifest_path.exists(): return { "manifest_version": 1, "split_count": 0, "splits": [], } return json.loads(manifest_path.read_text(encoding="utf-8")) def save_oracle_dataset_split_manifest(manifest: dict[str, Any], path: str | Path) -> None: target = Path(path) target.parent.mkdir(parents=True, exist_ok=True) target.write_text(json.dumps(manifest, sort_keys=True, indent=2) + "\n", encoding="utf-8") def upsert_oracle_dataset_split_manifest_entry( manifest_path: str | Path, *, split_dir: str | Path, summary: OracleDatasetSplitSummary, annotations: dict[str, Any] | None = None, ) -> dict[str, Any]: split_root = Path(split_dir) entry = OracleDatasetSplitManifestEntry( split_name=str(summary.split_name), split_dir=str(split_root), split_summary_path=str(split_root / "split_summary.json"), holdout_prompt_families=list(summary.holdout_prompt_families), holdout_prompt_variants=list(summary.holdout_prompt_variants), holdout_layers=list(summary.holdout_layers), train_label_count=int(summary.train_label_count), test_label_count=int(summary.test_label_count), annotations={} if annotations is None else dict(annotations), ) manifest_target = Path(manifest_path) manifest_target.parent.mkdir(parents=True, exist_ok=True) with manifest_target.open("a+", encoding="utf-8") as handle: fcntl.flock(handle.fileno(), fcntl.LOCK_EX) handle.seek(0) payload_text = handle.read() locked_manifest = ( { "manifest_version": 1, "split_count": 0, "splits": [], } if not payload_text.strip() else json.loads(payload_text) ) splits = [dict(item) for item in locked_manifest.get("splits", [])] updated = False for index, payload in enumerate(splits): if str(payload.get("split_name")) == entry.split_name or str(payload.get("split_dir")) == entry.split_dir: splits[index] = entry.to_dict() updated = True break if not updated: splits.append(entry.to_dict()) next_manifest = { **locked_manifest, "manifest_version": int(locked_manifest.get("manifest_version", 1)), "split_count": len(splits), "splits": splits, } handle.seek(0) handle.truncate() handle.write(json.dumps(next_manifest, sort_keys=True, indent=2) + "\n") handle.flush() fcntl.flock(handle.fileno(), fcntl.LOCK_UN) return next_manifest def _render_batch_summary_table(candidate_stats: dict[str, dict[str, Any]]) -> str: header = "| candidate | selected | safe_count | safe_rate | mean_total_bytes | mean_abs_error | mean_token_p95_error |" separator = "| --- | ---: | ---: | ---: | ---: | ---: | ---: |" rows = [header, separator] for candidate, stats in sorted( candidate_stats.items(), key=lambda item: (-int(item[1]["selected_count"]), -float(item[1]["safe_rate"]), float(item[1]["mean_total_bytes"]), item[0]), ): rows.append( "| " + " | ".join( [ candidate, str(int(stats["selected_count"])), str(int(stats["safe_count"])), f"{float(stats['safe_rate']):.3f}", f"{float(stats['mean_total_bytes']):.1f}", f"{float(stats['mean_abs_error']):.6f}", f"{float(stats['mean_token_p95_error']):.6f}", ] ) + " |" ) return "\n".join(rows)