"""Build a split-safe full-trajectory SFT dataset from completed runs.""" from __future__ import annotations import argparse from concurrent.futures import ThreadPoolExecutor, as_completed import hashlib import json import sys from collections import defaultdict from pathlib import Path from typing import Any, Dict, Iterable, List, Mapping, Sequence REPO_ROOT = Path(__file__).resolve().parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from eviscene.trajectory.schema import ( DatasetPerceptionExample, DatasetTrajectorySFTExample, PerceptionExample, TrajectorySFTExample, ) from eviscene.trajectory.exporter import export_trajectory_sft_example SPLITS = ("train", "validation", "test") ACCEPTED_RELEASE_SCHEMAS = frozenset( { "ifv-accepted-teacher-release-v2", "ifv-accepted-teacher-release-v3", "ifv-accepted-teacher-release-v4", } ) # Use bytes/3 for conservative admission; final token counts require # the actual model processor, especially near the context limit. UTF8_BYTES_PER_APPROX_TOKEN = 3 DEFAULT_SHORT_TRAJECTORY_MAX_TOKENS = 131072 DETERMINISTIC_FATAL_TEACHER_REASONS = frozenset( { "incorrect_result", "engineering_error", "legacy_core_ownership", } ) def _load_json(path: Path) -> Dict[str, Any]: value = json.loads(path.read_text(encoding="utf-8")) if not isinstance(value, dict): raise ValueError(f"expected JSON object: {path}") return value def _load_jsonl(path: Path) -> List[Dict[str, Any]]: if not path.is_file(): return [] rows: List[Dict[str, Any]] = [] for line_number, line in enumerate( path.read_text(encoding="utf-8").split("\n"), start=1, ): if not line.strip(): continue value = json.loads(line) if not isinstance(value, dict): raise ValueError(f"{path}:{line_number} must be a JSON object") rows.append(value) return rows def _write_json(path: Path, value: Mapping[str, Any]) -> None: path.parent.mkdir(parents=True, exist_ok=True) path.write_text( json.dumps(value, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) def _write_jsonl(path: Path, rows: Iterable[Mapping[str, Any]]) -> None: lines = [ json.dumps(row, ensure_ascii=False, separators=(",", ":")) for row in rows ] path.parent.mkdir(parents=True, exist_ok=True) path.write_text( "\n".join(lines) + ("\n" if lines else ""), encoding="utf-8", ) def _export_trace_file_job( job: tuple[str, Path, Path | None, Dict[str, Any]], ) -> tuple[str, TrajectorySFTExample | None, str]: """Export one trace while keeping raw trace data out of the coordinator.""" ( episode_id, trace_path, runtime_store_path, source_metadata, ) = job try: trace = _load_json(trace_path) if runtime_store_path is not None: state = dict(_mapping(trace.get("state"))) runtime_store = dict(_mapping(state.get("runtime_store"))) runtime_store["runtime_path"] = str(runtime_store_path) state["runtime_store"] = runtime_store trace["state"] = state return ( episode_id, export_trajectory_sft_example( trace, source_metadata=source_metadata, ), "", ) except Exception as exc: message = str(exc) or type(exc).__name__ if not isinstance(exc, ValueError): message = f"{type(exc).__name__}: {message}" return episode_id, None, message def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as handle: for chunk in iter(lambda: handle.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def _mapping(value: Any) -> Mapping[str, Any]: return value if isinstance(value, Mapping) else {} def _rows(value: Any) -> List[Mapping[str, Any]]: if not isinstance(value, list): return [] return [item for item in value if isinstance(item, Mapping)] def _collect_runtime_ids(trace: Mapping[str, Any]) -> set[str]: result: set[str] = set() def walk(value: Any) -> None: if isinstance(value, Mapping): for key, child in value.items(): if str(key).endswith("_id") and isinstance(child, str): if child.strip(): result.add(child.strip()) elif str(key).endswith("_ids") and isinstance(child, list): result.update( str(item).strip() for item in child if str(item).strip() ) walk(child) elif isinstance(value, list): for child in value: walk(child) walk(trace) return result def _cross_case_source_families(trace: Mapping[str, Any]) -> List[str]: """Return only source identities that can indicate cross-case leakage. ``source_family`` has two runtime meanings: a content digest when fetched content is available, and a registered-domain fallback otherwise. A shared content digest can make two cases near-duplicates; merely visiting the same host (for example Wikipedia) cannot. Domain families remain useful inside one investigation for source-diversity accounting, but must not merge unrelated cases into one train/validation component. """ state = _mapping(trace.get("state")) investigation = _mapping(state.get("investigation_state")) families = { str(item.get("source_family", "")).strip() for item in _rows(investigation.get("evidence")) if str(item.get("source_family", "")).strip().startswith("content:") } return sorted(families) def _trace_case_id(trace: Mapping[str, Any], episode_id: str) -> str: state = _mapping(trace.get("state")) runtime_case = _mapping(state.get("runtime_case")) return str( runtime_case.get("case_id") or trace.get("case_id") or trace.get("image_id") or episode_id ).strip() def _load_fixed_case_split(path: Path) -> Dict[str, Dict[str, Any]]: rows = _load_jsonl(path.expanduser().resolve()) result: Dict[str, Dict[str, Any]] = {} for row in rows: case_id = str(row.get("case_id", "")).strip() split = str(row.get("split", "")).strip() group_id = str(row.get("split_group_id", "")).strip() if not case_id or split not in {"train", "validation"} or not group_id: raise ValueError( "fixed case split rows require case_id, train/validation split, " "and split_group_id" ) if case_id in result: raise ValueError(f"duplicate fixed case split case_id: {case_id}") result[case_id] = dict(row) if not result: raise ValueError("fixed case split is empty") return result def _load_gate_artifacts( root: Path, *, suffix: str, ) -> Dict[str, Dict[str, Any]]: result: Dict[str, Dict[str, Any]] = {} for path in sorted(root.expanduser().resolve().rglob(f"*{suffix}")): payload = _load_json(path) episode_id = str( payload.get("episode_id") or _mapping(payload.get("rollout")).get("episode_id", "") ).strip() if not episode_id: raise ValueError(f"gate artifact lacks episode_id: {path}") if episode_id in result: raise ValueError(f"duplicate gate artifact episode_id: {episode_id}") result[episode_id] = payload return result def _load_canonical_accepted_release( path: Path, ) -> tuple[Path, Dict[str, Dict[str, Any]], Dict[str, Any]]: """Load the one canonical source of frozen teacher examples. The accepted release already contains the post-judge policy export. A frozen dataset export must consume those rows instead of reconstructing policy examples from the original rollout directories, whose historical exports may have used a stricter pre-judge verdict-chain gate. """ root = path.expanduser().resolve() manifest_path = root / "accepted_release_manifest.json" manifest = _load_json(manifest_path) if manifest.get("schema_version") not in ACCEPTED_RELEASE_SCHEMAS: raise ValueError( "accepted-release has unsupported schema: " f"{manifest.get('schema_version')!r}" ) if not (root / "run_manifest.json").is_file(): raise ValueError("accepted-release is missing run_manifest.json") eligibility_dir = root / "eligibility" if not eligibility_dir.is_dir(): raise ValueError("accepted-release is missing eligibility/") selected_rows = _load_jsonl(root / "selected_episodes.jsonl") selected_by_episode: Dict[str, Dict[str, Any]] = {} selected_cases: set[str] = set() for row in selected_rows: episode_id = str(row.get("episode_id", "")).strip() case_id = str(row.get("case_id", "")).strip() trace_sha256 = str(row.get("source_trace_sha256", "")).strip() if not episode_id or not case_id or not trace_sha256: raise ValueError( "accepted-release selected episode requires " "case_id, episode_id and source_trace_sha256" ) if episode_id in selected_by_episode: raise ValueError( f"accepted-release has duplicate episode_id: {episode_id}" ) if case_id in selected_cases: raise ValueError( f"accepted-release has duplicate selected case_id: {case_id}" ) if row.get("deterministic_hard_gate_pass") is not True: raise ValueError( "accepted-release contains a selected episode without a " f"passing deterministic gate: {episode_id}" ) selected_by_episode[episode_id] = row selected_cases.add(case_id) expected_count = int(manifest.get("accepted_case_count", -1)) if expected_count != len(selected_cases): raise ValueError( "accepted-release accepted_case_count does not match " f"selected episodes: {expected_count} != {len(selected_cases)}" ) if not selected_by_episode: raise ValueError("accepted-release selected_episodes.jsonl is empty") return root, selected_by_episode, manifest def _accepted_release_runtime_store( root: Path, row: Mapping[str, Any], *, required: bool, ) -> Path | None: raw = str(row.get("runtime_store_path", "")).strip() episode_id = str(row.get("episode_id", "")).strip() if not raw: if required: raise ValueError( "accepted-release v4 selected episode has no runtime store: " f"{episode_id}" ) return None root = root.resolve() runtime_store = (root / raw).resolve() if runtime_store != root and root not in runtime_store.parents: raise ValueError( "accepted-release runtime store escapes the release root: " f"{episode_id}" ) if not runtime_store.is_dir(): raise ValueError( "accepted-release runtime store is missing: " f"{episode_id}: {runtime_store}" ) if not any((runtime_store / "context").glob("*.json")): raise ValueError( "accepted-release runtime store has no context manifests: " f"{episode_id}" ) if not (runtime_store / "artifacts" / "sha256").is_dir(): raise ValueError( "accepted-release runtime store has no artifact store: " f"{episode_id}" ) expected_files = int(row.get("runtime_store_file_count", 0) or 0) if required and expected_files < 1: raise ValueError( "accepted-release v4 selected episode has invalid runtime store " f"metadata: {episode_id}" ) return runtime_store def _deterministic_teacher_quality(score: Mapping[str, Any]) -> Dict[str, Any]: if not score: return { "hard_gate_pass": False, "fatal_reasons": ["missing_training_quality_score"], "red_flags": [], } reasons = [ str(item) for item in score.get("training_exclusion_reasons", []) or [] if str(item).strip() ] fatal_reasons = [ reason for reason in reasons if reason in DETERMINISTIC_FATAL_TEACHER_REASONS ] if score.get("training_eligible") is False and not reasons: fatal_reasons.append("teacher_quality_gate_failed") return { "hard_gate_pass": not fatal_reasons, "fatal_reasons": list(dict.fromkeys(fatal_reasons)), "red_flags": [ reason for reason in reasons if reason not in DETERMINISTIC_FATAL_TEACHER_REASONS ], } class _UnionFind: def __init__(self, values: Iterable[str]) -> None: self.parent = {value: value for value in values} def find(self, value: str) -> str: parent = self.parent[value] if parent != value: self.parent[value] = self.find(parent) return self.parent[value] def union(self, left: str, right: str) -> None: root_left = self.find(left) root_right = self.find(right) if root_left != root_right: self.parent[max(root_left, root_right)] = min(root_left, root_right) def _assign_split( group_id: str, *, train_ratio: float, validation_ratio: float, seed: str, ) -> str: digest = hashlib.sha256(f"{seed}:{group_id}".encode("utf-8")).digest() value = int.from_bytes(digest[:8], "big") / 2**64 if value < train_ratio: return "train" if value < train_ratio + validation_ratio: return "validation" return "test" def export_dataset( run_dirs: Sequence[Path] | None, output_dir: Path, *, accepted_release_path: Path | None = None, case_split_path: Path | None = None, eligibility_dir: Path | None = None, semantic_reward_dir: Path | None = None, require_frozen_gates: bool = False, minimum_accepted_cases: int | None = None, require_all_validation_cases: bool = True, train_ratio: float = 0.8, validation_ratio: float = 0.1, seed: str = "ifv-policy-v1", short_max_tokens: int = DEFAULT_SHORT_TRAJECTORY_MAX_TOKENS, export_concurrency: int = 1, ) -> Dict[str, Any]: accepted_release: Path | None = None accepted_release_rows: Dict[str, Dict[str, Any]] = {} accepted_release_manifest: Dict[str, Any] = {} if accepted_release_path is not None: if run_dirs: raise ValueError( "accepted-release mode cannot be combined with --run-dir" ) ( accepted_release, accepted_release_rows, accepted_release_manifest, ) = ( _load_canonical_accepted_release(accepted_release_path) ) run_dirs = [accepted_release] canonical_eligibility_dir = accepted_release / "eligibility" if eligibility_dir is not None and ( eligibility_dir.expanduser().resolve() != canonical_eligibility_dir ): raise ValueError( "accepted-release mode owns its eligibility directory; " "do not pass a second --eligibility-dir" ) eligibility_dir = canonical_eligibility_dir require_frozen_gates = True if not run_dirs: raise ValueError("at least one run directory is required") frozen_inputs = (case_split_path, eligibility_dir, semantic_reward_dir) if any(value is not None for value in frozen_inputs): require_frozen_gates = True if require_frozen_gates and case_split_path is None: raise ValueError("frozen SFT export requires --case-split") if require_frozen_gates and eligibility_dir is None: raise ValueError("frozen SFT export requires --eligibility-dir") if minimum_accepted_cases is not None and minimum_accepted_cases < 0: raise ValueError("minimum_accepted_cases must be non-negative") if require_frozen_gates and minimum_accepted_cases is None: minimum_accepted_cases = 40 if train_ratio <= 0 or validation_ratio < 0: raise ValueError("split ratios must be non-negative") if train_ratio + validation_ratio >= 1: raise ValueError("train_ratio + validation_ratio must be below 1") if short_max_tokens < 1: raise ValueError("short_max_tokens must be positive") if export_concurrency < 1: raise ValueError("export_concurrency must be at least 1") if output_dir.exists() and any(output_dir.iterdir()): raise FileExistsError( f"dataset output directory must be new or empty: {output_dir}" ) fixed_split = ( _load_fixed_case_split(case_split_path) if case_split_path is not None else None ) eligibility_by_episode = ( _load_gate_artifacts(eligibility_dir, suffix=".sft_eligibility.json") if eligibility_dir is not None else {} ) semantic_by_episode = ( _load_gate_artifacts(semantic_reward_dir, suffix=".semantic_reward.json") if semantic_reward_dir is not None else {} ) trajectory_by_episode: Dict[str, TrajectorySFTExample] = {} perception_by_episode: Dict[str, PerceptionExample] = {} action_only_by_episode: Dict[str, Dict[str, Any]] = {} episode_metadata: Dict[str, Dict[str, Any]] = {} export_jobs: Dict[ str, tuple[str, Path, Path | None, Dict[str, Any]] ] = {} score_by_episode: Dict[str, Dict[str, Any]] = {} source_runs: List[Dict[str, Any]] = [] for raw_run_dir in run_dirs: run_dir = raw_run_dir.expanduser().resolve() manifest = _load_json(run_dir / "run_manifest.json") source_runs.append( { "run_id": manifest.get("run_id"), "path": str(run_dir), "git_commit": manifest.get("git_commit"), "release_id": _mapping(manifest.get("benchmark")).get( "release_id" ), } ) for row in _load_jsonl(run_dir / "trajectory_scores.jsonl"): score_key = str( row.get("episode_id") or row.get("case_id", "") ).strip() if score_key: score_by_episode[score_key] = dict(row) for row in _load_jsonl(run_dir / "perception_trajectories.jsonl"): example = PerceptionExample.model_validate(row) if example.episode_id in perception_by_episode: raise ValueError( "duplicate perception episode across source runs: " f"{example.episode_id}" ) perception_by_episode[example.episode_id] = example for trace_path in sorted((run_dir / "traces").glob("*.json")): trace = _load_json(trace_path) episode_id = str( trace.get("image_id") or trace_path.stem ).strip() case_id = _trace_case_id(trace, episode_id) trace_sha256 = _sha256(trace_path) score = score_by_episode.get(episode_id, {}) deterministic_quality = _deterministic_teacher_quality(score) accepted_release_row = ( accepted_release_rows.get(episode_id) if accepted_release is not None else None ) runtime_store_path: Path | None = None if accepted_release is not None: if accepted_release_row is None: raise ValueError( "accepted-release trace is not listed as selected: " f"{episode_id}" ) if str( accepted_release_row.get("source_trace_sha256", "") ) != trace_sha256: raise ValueError( "accepted-release trace SHA-256 mismatch: " f"{episode_id}" ) runtime_store_path = _accepted_release_runtime_store( accepted_release, accepted_release_row, required=( accepted_release_manifest.get("schema_version") == "ifv-accepted-teacher-release-v4" and str( accepted_release_row.get( "decision_policy_version", "", ) ) == "unified-react-v1" ), ) buckets = accepted_release_row.get( "training_buckets", ["reasoning_sft"], ) if ( isinstance(buckets, list) and "reasoning_sft" not in buckets ): # Action-only traces are intentionally excluded from the # reasoning policy corpus, but their independent # perception target remains valid training data. action_only_by_episode[episode_id] = { "case_id": case_id, "source_run_id": manifest.get("run_id"), "source_trace": str(trace_path), "runtime_store_path": ( str(runtime_store_path) if runtime_store_path is not None else None ), } continue score = { "episode_id": episode_id, "total": float( accepted_release_row.get("teacher_score", 0.0) or 0.0 ), "training_eligible": True, "training_exclusion_reasons": [], } deterministic_quality = { "hard_gate_pass": True, "fatal_reasons": [ str(item) for item in accepted_release_row.get( "deterministic_fatal_reasons", [] ) or [] ], "red_flags": [ str(item) for item in accepted_release_row.get( "deterministic_red_flags", [] ) or [] ], } training_eligible = bool(score.get("training_eligible", False)) eligibility = eligibility_by_episode.get(episode_id, {}) semantic = semantic_by_episode.get(episode_id, {}) if semantic: semantic_episode_id = str( _mapping(semantic.get("rollout")).get( "episode_id", semantic.get("episode_id", ""), ) ) if ( str(semantic.get("case_id", "")) != case_id or semantic_episode_id != episode_id or str( _mapping(semantic.get("source_trace")).get( "sha256", "" ) ) != trace_sha256 ): raise ValueError( "semantic diagnostic artifact does not match trace: " f"{episode_id}" ) if require_frozen_gates: eligibility_gates = _mapping(eligibility.get("gates")) split_row = (fixed_split or {}).get(case_id, {}) runtime_case = _mapping( _mapping(trace.get("state")).get("runtime_case") ) structured_gate_pass = bool( str(eligibility.get("case_id", "")) == case_id and str(eligibility.get("episode_id", "")) == episode_id and eligibility_gates.get("engineering_valid") is True and eligibility_gates.get("sft_eligibility_pass") is True and str( _mapping(eligibility.get("source_trace")).get( "sha256", "" ) ) == trace_sha256 and str(split_row.get("image_sha256", "")) == str(runtime_case.get("image_sha256", "")) ) training_eligible = bool( structured_gate_pass and deterministic_quality["hard_gate_pass"] ) exclusion_reasons = [ str(item) for item in score.get( "training_exclusion_reasons", ["missing_training_quality_score"], ) or [] ] if not score: exclusion_reasons = ["missing_training_quality_score"] if require_frozen_gates and not training_eligible: frozen_reasons = [] if not structured_gate_pass: frozen_reasons.append("structured_teacher_gate_failed") frozen_reasons.extend(deterministic_quality["fatal_reasons"]) exclusion_reasons = [ *frozen_reasons, *exclusion_reasons, ] elif require_frozen_gates: exclusion_reasons = [] episode_metadata[episode_id] = { "episode_id": episode_id, "case_id": case_id, "source_run_id": manifest.get("run_id"), "source_trace": str(trace_path), "source_family_keys": _cross_case_source_families(trace), "runtime_ids": sorted(_collect_runtime_ids(trace)), "teacher_score": float( score.get("total", 0.0) or 0.0 ), "training_eligible": training_eligible, "training_exclusion_reasons": exclusion_reasons, "deterministic_hard_gate_pass": deterministic_quality[ "hard_gate_pass" ], "deterministic_fatal_reasons": deterministic_quality[ "fatal_reasons" ], "deterministic_red_flags": deterministic_quality["red_flags"], "structured_eligibility_pass": eligibility.get("gates", {}).get( "sft_eligibility_pass" ), "semantic_audit_pass": semantic.get("gates", {}).get( "semantic_audit_pass" ), "sft_eligibility_artifact_id": eligibility.get("artifact_id"), "semantic_reward_artifact_id": semantic.get("artifact_id"), "decision_policy_version": str( trace.get("decision_policy_version") or _mapping(trace.get("state")).get( "decision_policy_version", "", ) or "" ), "runtime_store_path": ( str(runtime_store_path) if runtime_store_path is not None else None ), } export_jobs[episode_id] = ( episode_id, trace_path, runtime_store_path, { "source_run_id": manifest.get("run_id"), "runtime_commit": manifest.get("git_commit"), "release_id": _mapping(manifest.get("benchmark")).get( "release_id" ), "runtime_contract_version": _mapping( manifest.get("benchmark") ).get("runtime_contract_version", ""), "process_reference_protocol_version": manifest.get( "process_reference_protocol_version", "", ), }, ) ordered_jobs = [ export_jobs[episode_id] for episode_id in sorted(export_jobs) ] executor: ThreadPoolExecutor | None = None if export_concurrency == 1 or len(ordered_jobs) <= 1: export_results = ( _export_trace_file_job(job) for job in ordered_jobs ) for episode_id, trajectory, export_error in export_results: metadata = episode_metadata[episode_id] metadata["trajectory_export_error"] = export_error if trajectory is None: continue trajectory_by_episode[episode_id] = trajectory metadata["trajectory_token_count_estimate"] = ( trajectory.token_count_estimate ) metadata["tool_call_count"] = trajectory.tool_call_count else: executor = ThreadPoolExecutor( max_workers=min(export_concurrency, len(ordered_jobs)), thread_name_prefix="trajectory-export", ) futures = { executor.submit(_export_trace_file_job, job): job[0] for job in ordered_jobs } try: for future in as_completed(futures): episode_id, trajectory, export_error = future.result() metadata = episode_metadata[episode_id] metadata["trajectory_export_error"] = export_error if trajectory is None: continue trajectory_by_episode[episode_id] = trajectory metadata["trajectory_token_count_estimate"] = ( trajectory.token_count_estimate ) metadata["tool_call_count"] = trajectory.tool_call_count finally: executor.shutdown(wait=True) all_episodes = sorted(episode_metadata) policy_versions = { str(metadata["decision_policy_version"]) for metadata in episode_metadata.values() } if len(policy_versions) > 1: raise ValueError( "SFT dataset cannot mix decision_policy_version values: " + ", ".join(sorted(policy_versions)) ) decision_policy_version = next(iter(policy_versions), "") if require_frozen_gates: missing_split = sorted( { str(metadata["case_id"]) for metadata in episode_metadata.values() if str(metadata["case_id"]) not in (fixed_split or {}) } ) if missing_split: raise ValueError( "fixed case split lacks teacher cases: " + ", ".join(missing_split[:10]) ) if require_frozen_gates: for episode_id in all_episodes: metadata = episode_metadata[episode_id] if episode_id not in trajectory_by_episode: metadata["training_eligible"] = False metadata["training_exclusion_reasons"] = list( dict.fromkeys( [ "no_exportable_trajectory_sft", *metadata["training_exclusion_reasons"], ] ) ) eligible_by_case: Dict[str, List[str]] = defaultdict(list) for episode_id in all_episodes: metadata = episode_metadata[episode_id] if metadata["training_eligible"]: eligible_by_case[str(metadata["case_id"])].append(episode_id) selected_by_case: Dict[str, str] = {} for case_id, candidates in eligible_by_case.items(): selected_by_case[case_id] = max( candidates, key=lambda episode_id: ( -len(episode_metadata[episode_id]["deterministic_red_flags"]), float(episode_metadata[episode_id]["teacher_score"]), episode_id, ), ) for episode_id in all_episodes: metadata = episode_metadata[episode_id] if ( metadata["training_eligible"] and selected_by_case.get(str(metadata["case_id"])) != episode_id ): metadata["training_eligible"] = False metadata["training_exclusion_reasons"] = list( dict.fromkeys( [ "lower_quality_rollout_for_case", *metadata["training_exclusion_reasons"], ] ) ) episodes = sorted(selected_by_case.values()) accepted_case_ids = { str(episode_metadata[episode_id]["case_id"]) for episode_id in episodes } if minimum_accepted_cases is not None and ( len(accepted_case_ids) < minimum_accepted_cases ): raise ValueError( "frozen SFT export accepted too few cases: " f"{len(accepted_case_ids)} < {minimum_accepted_cases}" ) if require_all_validation_cases: policy_case_ids = { str(episode_metadata[episode_id]["case_id"]) for episode_id in all_episodes } expected_validation = { case_id for case_id, row in (fixed_split or {}).items() if ( row.get("split") == "validation" and case_id in policy_case_ids ) } missing_validation = sorted(expected_validation - accepted_case_ids) if missing_validation: raise ValueError( "frozen SFT export lacks accepted validation cases: " + ", ".join(missing_validation[:10]) ) else: episodes = [ episode_id for episode_id in all_episodes if episode_metadata[episode_id]["training_eligible"] and episode_id in trajectory_by_episode ] excluded_episode_rows = [ { **episode_metadata[episode_id], "trajectory_sft_rows": int(episode_id in trajectory_by_episode), "perception_example_count": int( episode_id in perception_by_episode ), } for episode_id in all_episodes if episode_id not in episodes ] union = _UnionFind(episodes) family_episodes: Dict[str, List[str]] = defaultdict(list) for episode_id in episodes: for family in episode_metadata[episode_id]["source_family_keys"]: family_episodes[family].append(episode_id) for members in family_episodes.values(): for member in members[1:]: union.union(members[0], member) components: Dict[str, List[str]] = defaultdict(list) for episode_id in episodes: components[union.find(episode_id)].append(episode_id) split_by_episode: Dict[str, str] = {} group_by_episode: Dict[str, str] = {} for members in components.values(): if fixed_split is not None: member_case_splits = { str(fixed_split[str(episode_metadata[item]["case_id"])]["split"]) for item in members } if len(member_case_splits) != 1: raise ValueError( "fixed case split separates one source-family group: " + ", ".join(sorted(member_case_splits)) ) for episode_id in members: fixed_row = fixed_split[ str(episode_metadata[episode_id]["case_id"]) ] split_by_episode[episode_id] = str(fixed_row["split"]) group_by_episode[episode_id] = str( fixed_row["split_group_id"] ) continue group_payload = { "episodes": sorted(members), "families": sorted( { family for episode_id in members for family in episode_metadata[episode_id][ "source_family_keys" ] } ), } group_id = hashlib.sha256( json.dumps( group_payload, sort_keys=True, separators=(",", ":"), ).encode("utf-8") ).hexdigest()[:20] split = _assign_split( group_id, train_ratio=train_ratio, validation_ratio=validation_ratio, seed=seed, ) for episode_id in members: split_by_episode[episode_id] = split group_by_episode[episode_id] = group_id split_rows: Dict[str, List[Dict[str, Any]]] = { split: [] for split in SPLITS } perception_split_rows: Dict[str, List[Dict[str, Any]]] = { split: [] for split in SPLITS } long_holdout_rows: List[Dict[str, Any]] = [] metadata_rows: List[Dict[str, Any]] = [] accepted_episode_rows: List[Dict[str, Any]] = [] for episode_id in episodes: metadata = episode_metadata[episode_id] split = split_by_episode[episode_id] trajectory = trajectory_by_episode[episode_id] approximate_tokens = ( trajectory.token_count_estimate + UTF8_BYTES_PER_APPROX_TOKEN - 1 ) // UTF8_BYTES_PER_APPROX_TOKEN is_short = approximate_tokens <= short_max_tokens accepted_row = { **metadata, "split": split, "split_group_id": group_by_episode[episode_id], "sft_route": "short_sft" if is_short else "long_holdout", "short_max_tokens": short_max_tokens, "trajectory_approx_tokens": approximate_tokens, } metadata_rows.append(accepted_row) accepted_episode_rows.append( { key: accepted_row.get(key) for key in ( "case_id", "episode_id", "split", "split_group_id", "source_run_id", "source_trace", "sft_eligibility_artifact_id", "semantic_reward_artifact_id", "deterministic_hard_gate_pass", "deterministic_fatal_reasons", "deterministic_red_flags", "teacher_score", "trajectory_token_count_estimate", "tool_call_count", "sft_route", ) } ) dataset_trajectory = DatasetTrajectorySFTExample( **trajectory.model_dump(mode="json"), dataset_version=( "ifv-trajectory-sft-dataset-v3" if decision_policy_version == "unified-react-v1" else "ifv-trajectory-sft-dataset-v2" ), split=split, split_group_id=group_by_episode[episode_id], source_family_keys=metadata["source_family_keys"], teacher_score=metadata["teacher_score"], ).model_dump(mode="json") if is_short: split_rows[split].append(dataset_trajectory) else: long_holdout_rows.append( { **dataset_trajectory, "holdout_reason": "trajectory_exceeds_short_token_budget", "short_max_tokens": short_max_tokens, } ) perception_example = perception_by_episode.get(episode_id) if perception_example is not None and is_short: dataset_perception = DatasetPerceptionExample( **perception_example.model_dump(mode="json"), split=split, split_group_id=group_by_episode[episode_id], ) perception_split_rows[split].append( dataset_perception.model_dump(mode="json") ) # Perception is an independent image-only task. Do not make it depend on # the policy trace having readable provider thought. In particular, an # accepted action-only trace is excluded from reasoning SFT but still has # a valid PerceptionReport target. action_only_perception_count = 0 for episode_id, action_only in action_only_by_episode.items(): perception_example = perception_by_episode.get(episode_id) if perception_example is None: continue case_id = str(action_only["case_id"]) if fixed_split is not None: fixed_row = fixed_split.get(case_id) if fixed_row is None: raise ValueError( "fixed case split lacks action-only perception case: " f"{case_id}" ) split = str(fixed_row["split"]) split_group_id = str(fixed_row["split_group_id"]) else: # Canonical accepted-release exports require a frozen split. Keep # this invariant explicit instead of inventing another split. raise ValueError( "action-only perception export requires a frozen case split" ) dataset_perception = DatasetPerceptionExample( **perception_example.model_dump(mode="json"), split=split, split_group_id=split_group_id, ) perception_split_rows[split].append( dataset_perception.model_dump(mode="json") ) action_only_perception_count += 1 for split, rows in split_rows.items(): rows.sort(key=lambda item: item["episode_id"]) _write_jsonl(output_dir / f"{split}.jsonl", rows) perception_rows = perception_split_rows[split] perception_rows.sort(key=lambda item: item["episode_id"]) _write_jsonl( output_dir / f"perception.{split}.jsonl", perception_rows, ) long_holdout_rows.sort(key=lambda item: item["episode_id"]) _write_jsonl(output_dir / "long_holdout.jsonl", long_holdout_rows) metadata_rows.sort(key=lambda item: item["episode_id"]) _write_jsonl(output_dir / "episode_metadata.jsonl", metadata_rows) accepted_episode_rows.sort(key=lambda item: item["episode_id"]) _write_jsonl(output_dir / "accepted_episodes.jsonl", accepted_episode_rows) excluded_episode_rows.sort(key=lambda item: item["episode_id"]) _write_jsonl( output_dir / "excluded_episode_metadata.jsonl", excluded_episode_rows, ) action_only_rows: List[Dict[str, Any]] = [] if accepted_release is not None: action_only_rows = _load_jsonl(accepted_release / "action_only.jsonl") _write_jsonl(output_dir / "action_only.jsonl", action_only_rows) selected_release_case_count = ( len(accepted_release_rows) if accepted_release is not None else len(episodes) ) manifest = { "schema_version": "ifv-trajectory-sft-dataset-manifest-v1", "dataset_version": ( "ifv-trajectory-sft-dataset-v3" if decision_policy_version == "unified-react-v1" else "ifv-trajectory-sft-dataset-v2" ), "trajectory_version": ( "ifv-trajectory-sft-v3" ), "decision_policy_version": decision_policy_version, "perception_version": "ifv-perception-v1", "legacy_step_policy_export": "disabled", "split_mode": ( "frozen_teacher_sft" if require_frozen_gates else "derived" ), "seed": seed, "short_max_tokens": short_max_tokens, "short_max_tokens_unit": "approximate_model_tokens", "token_estimate_unit": "utf8_bytes", "split_ratios": ( None if require_frozen_gates else { "train": train_ratio, "validation": validation_ratio, "test": 1 - train_ratio - validation_ratio, } ), "source_runs": source_runs, "frozen_inputs": ( { "case_split": { "path": str(case_split_path.expanduser().resolve()), "sha256": _sha256(case_split_path.expanduser().resolve()), }, "eligibility_dir": str(eligibility_dir.expanduser().resolve()), "semantic_reward_dir": str( semantic_reward_dir.expanduser().resolve() ) if semantic_reward_dir is not None else None, "accepted_release": ( { "path": str(accepted_release), "manifest_sha256": _sha256( accepted_release / "accepted_release_manifest.json" ), "schema_version": accepted_release_manifest.get( "schema_version" ), "runtime_store_archive": accepted_release_manifest.get( "runtime_store_archive" ), } if accepted_release is not None else None ), } if require_frozen_gates else None ), "source_mode": ( "canonical_accepted_release" if accepted_release is not None else "rollout_run" ), "export_concurrency": export_concurrency, "episode_count": sum(len(rows) for rows in split_rows.values()), "accepted_episode_count": len(episodes), "selected_release_case_count": selected_release_case_count, "action_only_episode_count": len(action_only_rows), "action_only_perception_count": action_only_perception_count, "long_holdout_episode_count": len(long_holdout_rows), "accepted_case_count": len( {str(episode_metadata[episode_id]["case_id"]) for episode_id in episodes} ), "short_sft_case_count": len( { str(row["case_id"]) for rows in split_rows.values() for row in rows } ), "candidate_episode_count": len(all_episodes), "excluded_episode_count": len(excluded_episode_rows), "group_count": len(components), "example_counts": { split: len(rows) for split, rows in split_rows.items() }, "perception_example_counts": { split: len(rows) for split, rows in perception_split_rows.items() }, "artifacts": { "train": "train.jsonl", "validation": "validation.jsonl", "test": "test.jsonl", "perception_train": "perception.train.jsonl", "perception_validation": "perception.validation.jsonl", "perception_test": "perception.test.jsonl", "episode_metadata": "episode_metadata.jsonl", "accepted_episodes": "accepted_episodes.jsonl", "excluded_episode_metadata": ( "excluded_episode_metadata.jsonl" ), "long_holdout": "long_holdout.jsonl", "action_only": "action_only.jsonl", }, } _write_json(output_dir / "manifest.json", manifest) return manifest def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--run-dir", action="append", type=Path) parser.add_argument("--accepted-release", type=Path) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--case-split", type=Path) parser.add_argument("--eligibility-dir", type=Path) parser.add_argument("--semantic-reward-dir", type=Path) parser.add_argument("--minimum-accepted-cases", type=int) parser.add_argument("--allow-missing-validation-cases", action="store_true") parser.add_argument("--train-ratio", type=float, default=0.8) parser.add_argument("--validation-ratio", type=float, default=0.1) parser.add_argument("--seed", default="ifv-policy-v1") parser.add_argument( "--short-max-tokens", type=int, default=DEFAULT_SHORT_TRAJECTORY_MAX_TOKENS, help=( "Maximum estimated tokens for short full-trajectory SFT rows; " "longer accepted episodes are kept in long_holdout.jsonl." ), ) parser.add_argument( "--export-concurrency", type=int, default=1, help=( "number of concurrent trace-to-SFT exporter workers; output " "ordering remains stable by episode_id" ), ) return parser def main() -> None: args = build_parser().parse_args() if bool(args.run_dir) == bool(args.accepted_release): raise SystemExit( "provide exactly one of --run-dir or --accepted-release" ) result = export_dataset( args.run_dir, args.output_dir, accepted_release_path=args.accepted_release, case_split_path=args.case_split, eligibility_dir=args.eligibility_dir, semantic_reward_dir=args.semantic_reward_dir, minimum_accepted_cases=args.minimum_accepted_cases, require_all_validation_cases=not args.allow_missing_validation_cases, train_ratio=args.train_ratio, validation_ratio=args.validation_ratio, seed=args.seed, short_max_tokens=args.short_max_tokens, export_concurrency=args.export_concurrency, ) print(json.dumps(result, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()