Skip to content

The recording format

The recording is the contract between the recorder and the doctor: the recorder knows nothing about diagnosis, the doctor knows nothing about trainers. One directory per run, self-contained, rsync-able, committable to blob storage:

run_dir/
  manifest.json          run metadata: config, environment, git sha, versions
  steps.parquet          per-step scalar metrics (one wide table)
  samples/
    shard-00000.parquet  per-sample forensics (sharded every 4096 samples)
    shard-00001.parquet
  events.jsonl           sparse discrete events, append-only
  rank-001/              optional: per-rank sub-recording (distributed, opt-in)
    manifest.json ...

Typed access goes through RunReader (and RunView for numpy series); the schemas are pydantic models in rlprobe.schemas. The manifest carries schema_version; a reader refuses recordings written by a newer schema than it understands, so a version bump can never be silently misread.

Schema versioning policy

The recording format is versioned independently of the package (schema_version in every manifest; currently 1), with these rules:

  • Additive changes do not bump the version. New optional fields are the normal way the format evolves: every field except step is already optional, old recordings read as None for fields that didn't exist yet (detectors treat missing as "stay silent"), and old readers ignore columns they don't know. Adding a field is forward- and backward-compatible by construction.
  • Breaking changes bump the version — renaming/retyping a field, changing the meaning of an existing one, or restructuring the directory layout. A newer reader remains responsible for reading all older versions (migrating on read); an older reader refuses newer recordings with an explicit "upgrade rlprobe" error rather than misreading them.
  • Recordings are never rewritten in place. A recording is evidence; any future migration tooling produces a new run directory and leaves the original intact.

Practical consequence: pip install -U rlprobe can never orphan your archived runs, and an old pinned rlprobe fails loudly — never silently wrong — on runs recorded by a newer one.

manifest.json — RunManifest

Written once at run start: run_id, created_at, algorithm (grpo/ppo/…), model_name, world_size, the full trainer_config dump, and environment (git sha, package-version subset, GPU info when torch is present — all captured automatically unless capture_env=False).

steps.parquet — StepRecord

One row per optimizer step. Everything except step is optional: different algorithms and integration depths populate different subsets, and detectors must tolerate missing columns (a detector whose input field is absent simply stays silent).

group fields
reward reward_mean/std/min/max, reward_components (per-component breakdown)
KL to reference kl_mean, kl_max, kl_positional (coarse token-position buckets)
policy entropy_mean, logprob_mean, ratio_mean, ratio_p95, clip_fraction
advantages adv_mean, adv_std, adv_zero_frac
GRPO group_reward_std_mean, zero_variance_group_frac
PPO (value head) value_loss, value_explained_variance, value_clip_fraction
optimization loss, lr, grad_norm, grad_norm_post_clip, nan_count, inf_count, weight_update_ratio
generation completion_len_mean/p95/max, max_len_hit_frac, eos_rate, repetition_rate, truncation_rate
throughput step_time_s, generation_time_s, optimization_time_s
open-ended extra (any numeric metric you log that isn't a known field)

Dict-valued fields (reward_components, extra) are flattened to prefixed parquet columns (reward_component.<name>, extra.<name>) so the table stays wide and directly queryable by any parquet tool; the reader reconstructs the dicts.

samples/ — SampleRecord

Per-sample forensics: prompt_text, completion_text (or completion_token_ids), reward_total, per-component reward_components, advantage, logprob_mean, kl_contribution, completion_len, truncated, GRPO group_id, and capture_reason.

Capturing every rollout would dwarf the metrics, so the recorder applies a stratified sampling policy per log_samples call: keep the top-k and bottom-k by reward (the extremes are where hacking and reward bugs live) plus a random draw from the middle (so the typical distribution stays representable). Defaults: 2 + 2 + 4 per step, configurable via SamplingConfig; pass select=False to keep everything. Each kept sample is tagged with why (top_k / bottom_k / random / all).

events.jsonl — Event

Sparse, append-only, written to disk immediately (events must survive crashes): timestamp, kind, step, rank, message, data. Known kinds include train_begin, train_end, checkpoint, nan_detected, crash, auto_halt, recorder_drops, bad_record, shim_unavailable — but kind is an open string, so integrations can add their own.

Crash safety (the black box survives the crash)

The Recorder's log_* calls only enqueue onto a bounded queue; a daemon writer thread does all validation, serialization, and IO. Design consequences, all deliberate:

  • Training is never blocked. If the queue fills, records are dropped and counted (surfaced as a recorder_drops event) rather than back-pressuring the training loop.
  • Invalid values never raise at the call site. A malformed record is dropped and a bad_record event is written instead — a bad metric must not kill either the training loop or the writer thread mid-run.
  • Incremental flushing. During a live run, steps go to immutable steps/part-*.parquet files (every 5s or 200 steps); a clean close compacts them into the single steps.parquet. The reader transparently handles either layout, so crashed and still-running runs are fully readable — this is how live diagnosis and post-mortems on dead runs work.
  • Exit paths are covered: atexit, a sys.excepthook that logs a crash event with the traceback before chaining, and a SIGTERM handler. A hard kill (SIGKILL-class) loses at most the unflushed window.
  • The first parquet write is pre-warmed at writer-thread start (codec/DLL init can cost seconds on some machines) so a crash in a run's first seconds can't lose everything to a cold codec cache.

Measured overhead on a mock training loop is under 2% of step time; the benchmark (benchmarks/overhead.py) runs in CI. At 100k steps a recording is ~40 MB and every read surface (doctor, report, diff, watcher) stays in single-digit seconds (benchmarks/scale.py).

Distributed runs

Default: rank 0 records, every other rank is a no-op. Recorder(rank=N) for N ≠ 0 constructs an inert recorder — nothing written, no handlers installed, every method safe to call — so you can construct one per process with no if rank == 0: gating. The rationale: HF/TRL all-reduce logged scalars before callbacks fire, so rank 0 already sees global aggregates. Per-rank step rows would be duplicates — and would corrupt detector series, which must never see the same step twice.

Opt-in: record_all_ranks=True. Each non-zero rank writes a complete sub-recording under run_dir/rank-NNN/. The parent RunReader merges samples and events from rank dirs (rank-attributed, events sorted by time) while steps() stays rank-0-only — scalars remain the single authoritative series. Because every consumer (doctor, report, diff, forensics, watcher) reads through RunReader, the merge propagates everywhere automatically.

When to use which: for TRL, the default is complete — TRL gathers its per-sample log buffer to every process, so rank 0 already records the full global batch (verified empirically; record_all_ranks would write pure duplicates). The opt-in exists for manual-API training loops that log rank-local data, where per-rank forensics are real information.

One rank must also mean one guardian: the TRL callback runs live diagnosis and auto-halt only on the recording rank, so a shared recording can't make every worker re-alert on the same findings.