From 8448ada000c344218ede35747f69b4d05cae4d8c Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Thu, 6 Aug 2026 11:47:00 +0200 Subject: [PATCH 01/12] Init commit --- tools/merge_batches.py | 38 ++++++++++ tools/sync_loop.sh | 162 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 200 insertions(+) create mode 100644 tools/merge_batches.py create mode 100755 tools/sync_loop.sh diff --git a/tools/merge_batches.py b/tools/merge_batches.py new file mode 100644 index 0000000..abbd770 --- /dev/null +++ b/tools/merge_batches.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python +"""Merge the batch_* LeRobot dataset dirs produced by sync_loop.sh into one dataset.""" + +import argparse +import sys +from pathlib import Path + +from lerobot.datasets.aggregate import aggregate_datasets + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--batches-root", required=True, type=Path, help="Dir containing batch_* dataset dirs") + parser.add_argument("--output", required=True, type=Path, help="Output dir for the merged dataset") + parser.add_argument("--repo-id", default="merged", help="Identifier for the merged dataset") + args = parser.parse_args() + + batches = sorted(p for p in args.batches_root.glob("batch_*") if p.is_dir()) + if not batches: + print(f"No batch_* dirs found under {args.batches_root}", file=sys.stderr) + sys.exit(1) + + if args.output.exists(): + print(f"Output dir already exists, refusing to overwrite: {args.output}", file=sys.stderr) + sys.exit(1) + + print(f"Merging {len(batches)} batches into {args.output}") + aggregate_datasets( + repo_ids=[b.name for b in batches], + aggr_repo_id=args.repo_id, + roots=batches, + aggr_root=args.output, + ) + print("Merge complete.") + + +if __name__ == "__main__": + main() diff --git a/tools/sync_loop.sh b/tools/sync_loop.sh new file mode 100755 index 0000000..9b03da9 --- /dev/null +++ b/tools/sync_loop.sh @@ -0,0 +1,162 @@ +#!/usr/bin/env bash +# Pull recordings from the WS while collection is still running there, and +# export them in the background as they arrive, so pulling + exporting happen +# in parallel instead of waiting for all 1000 episodes then exporting. +# +# A recording is only "done" once /recording.rrd exists (it's +# written once on stop; chunks/ alone means still recording) — that's the +# completeness check, both for what's safe to export. +# +# nova-data-cli has no incremental/append mode (--output must not exist), so +# this can't grow one dataset live: each newly-arrived batch of recordings +# gets exported into its own batch_N output dir. Once collection is finished +# (idle timeout below), the batches are merged into one final LeRobot dataset +# via lerobot's own aggregate_datasets (tools/merge_batches.py). +set -euo pipefail +shopt -s nullglob + +REMOTE_HOST="intern@172.31.11.129" +REMOTE_DIRS=( + "/mnt/data/sebastian/raw_datasets/pick_and_place_sim_20260805_191245" +) +LOCAL_DEST="/home/sebi/ws/Data/raw_data/choreo2" + +NOVA_CLI_DIR="/home/sebi/ws/nova-data-cli" +EXPORT_CONFIG="/home/sebi/ws/pick_and_place_imitation_learning/data_collection/configs/lerobot_export.json" +EXPORT_OUTPUT_ROOT="/home/sebi/ws/Data/choreo2_export" + +LEDGER="${LOCAL_DEST}/.exported_ids" +INFLIGHT="${LOCAL_DEST}/.inflight_ids" +EXPORT_PID_FILE="${LOCAL_DEST}/.export.pid" + +POLL_SECONDS=180 +IDLE_MINUTES=10 + +mkdir -p "$LOCAL_DEST" "$EXPORT_OUTPUT_ROOT" +touch "$LEDGER" "$INFLIGHT" + +idle_polls_needed=$(( (IDLE_MINUTES * 60 + POLL_SECONDS - 1) / POLL_SECONDS )) +idle_count=0 +last_mtime="" +batch_num=0 + +remote_max_mtime() { + ssh "$REMOTE_HOST" "find ${REMOTE_DIRS[*]} -type f -printf '%T@\n' 2>/dev/null | sort -n | tail -1" +} + +sync_once() { + rsync -avP "${REMOTE_DIRS[@]/#/$REMOTE_HOST:}" "$LOCAL_DEST/" +} + +export_running() { + [[ -f "$EXPORT_PID_FILE" ]] && kill -0 "$(cat "$EXPORT_PID_FILE")" 2>/dev/null +} + +# Recordings that have a recording.rrd (complete) and aren't already +# exported or queued in a currently-running batch. +find_new_recordings() { + for remote_dir in "${REMOTE_DIRS[@]}"; do + local_root="${LOCAL_DEST}/$(basename "$remote_dir")" + for rrd in "$local_root"/*/recording.rrd; do + recording_dir="$(dirname "$rrd")" + recording_id="$(basename "$recording_dir")" + if ! grep -qxF "$recording_id" "$LEDGER" && ! grep -qxF "$recording_id" "$INFLIGHT"; then + echo "$recording_dir" + fi + done + done +} + +run_export_batch() { + local batch_dir="$1" + shift + local recordings=("$@") + + local merged_dir="${EXPORT_OUTPUT_ROOT}/_merged_${batch_dir}" + local output_dir="${EXPORT_OUTPUT_ROOT}/${batch_dir}" + + rm -rf "$merged_dir" + mkdir -p "$merged_dir" + for recording in "${recordings[@]}"; do + ln -s "$recording" "${merged_dir}/$(basename "$recording")" + done + + if (cd "$NOVA_CLI_DIR" && uv run nova-data-cli \ + --dataset "$merged_dir" \ + --config "$EXPORT_CONFIG" \ + --output "$output_dir"); then + for recording in "${recordings[@]}"; do + basename "$recording" >> "$LEDGER" + done + echo "Exported ${#recordings[@]} recordings -> $output_dir" + else + echo "Export batch $batch_dir failed, will retry these next round: ${recordings[*]}" >&2 + fi + + for recording in "${recordings[@]}"; do + sed -i "\|^$(basename "$recording")\$|d" "$INFLIGHT" + done + rm -rf "$merged_dir" +} + +maybe_start_export_batch() { + if export_running; then + return + fi + + new_recordings=() + while IFS= read -r line; do + [[ -n "$line" ]] && new_recordings+=("$line") + done < <(find_new_recordings) + + if [[ ${#new_recordings[@]} -eq 0 ]]; then + return + fi + + batch_num=$((batch_num + 1)) + batch_name="batch_$(printf '%04d' "$batch_num")" + for recording in "${new_recordings[@]}"; do + basename "$recording" >> "$INFLIGHT" + done + + echo "Starting export batch $batch_name with ${#new_recordings[@]} new recordings" + run_export_batch "$batch_name" "${new_recordings[@]}" & + echo $! > "$EXPORT_PID_FILE" +} + +while true; do + sync_once + maybe_start_export_batch + + mtime="$(remote_max_mtime)" + if [[ "$mtime" == "$last_mtime" ]]; then + idle_count=$((idle_count + 1)) + else + idle_count=0 + last_mtime="$mtime" + fi + + if [[ $idle_count -ge $idle_polls_needed ]]; then + echo "No new remote data for ${IDLE_MINUTES}m, assuming collection finished." + break + fi + + sleep "$POLL_SECONDS" +done + +echo "Final sync + export of any remaining recordings..." +sync_once +while export_running; do + sleep 5 +done +maybe_start_export_batch +while export_running; do + sleep 5 +done + +MERGED_OUTPUT="${EXPORT_OUTPUT_ROOT}_merged" +echo "All batches exported. Merging into $MERGED_OUTPUT..." +(cd "$NOVA_CLI_DIR" && uv run python tools/merge_batches.py \ + --batches-root "$EXPORT_OUTPUT_ROOT" \ + --output "$MERGED_OUTPUT") +echo "Done. Merged dataset: $MERGED_OUTPUT" From d37d7aff7e36448ae2ff88d7493aaef113c36008 Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Thu, 6 Aug 2026 12:06:16 +0200 Subject: [PATCH 02/12] Doc: updated docs --- README.md | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/README.md b/README.md index 8b098a1..715b2d7 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,30 @@ cameras. Examples: [`examples/`](examples/). Schema: config field, the export formats, camera resizing, and how the trimming modes choose episode boundaries (with figures). +## Live sync + export (`tools/sync_loop.sh`) + +Pulls recordings from a remote machine while collection is still running there, +and exports them in the background as each one completes — instead of waiting +for collection to finish before exporting anything. + +- **Two machines (collector + this one):** requires passwordless SSH to the + remote host (`ssh-copy-id`), since the script polls it every `POLL_SECONDS` + via `rsync`/`ssh`. Edit `REMOTE_HOST`, `REMOTE_DIRS`, and `LOCAL_DEST` at the + top of the script first. +- **Same machine:** SSH isn't needed if collection and export run on one box — + point `REMOTE_DIRS`/`LOCAL_DEST` at local paths and swap `sync_once`'s + `rsync` for a local copy (or skip syncing and export straight from the + collection dir). Not built yet; the script currently assumes a remote host. +- The script stops polling once the remote dir has been idle for + `IDLE_MINUTES`, does one final sync + export pass, then merges all batches + into one LeRobot dataset via `tools/merge_batches.py` (nova-data-cli itself + has no incremental/append mode, so each batch is a separate `--output` dir + until merged). + +```bash +tools/sync_loop.sh +``` + ## Tests ```bash From a8a5ddca42d000b6e573f9c31c344665ab767008 Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Thu, 6 Aug 2026 17:30:51 +0200 Subject: [PATCH 03/12] Replace sync_loop.sh with a parallel claim/commit export pipeline nova-data-cli has no incremental-append mode and is single-threaded per invocation, so exporting ~1000 episodes sequentially after collection wastes hours. pipeline.sh pulls (remote SSH/rsync or a local dir) and exports in parallel across workers sized from the machine's own RAM/cores, while collection is still running. Workers dynamically claim recordings via atomic mkdir, commit via a single atomic rename that doubles as the durable "done" record (closing a duplicate- export race found and fixed during testing), and bisect a failing batch instead of quarantining healthy batch-mates alongside a bad recording. Per-worker memory is capped via a systemd-run cgroup (MemorySwapMax=0) after an OOM-killed rerun subprocess was found to be the cause of an earlier corrupted export. merge_batches.py gains an atomic-output guarantee and a pre-merge video-encoder compatibility check that lerobot's own aggregate_datasets skips. Full design rationale in tools/AGENT.md. Test suite under tools/tests/ exercises both acquisition modes, live-feed and backlog-only, crash/restart, and a real end-to-end smoke test. Co-Authored-By: Claude Sonnet 5 --- README.md | 48 +-- src/nova_export/cli.py | 19 ++ tools/AGENT.md | 176 +++++++++++ tools/merge_batches.py | 91 +++++- tools/pipeline.sh | 392 +++++++++++++++++++++++++ tools/sync_loop.sh | 162 ---------- tools/tests/README.md | 100 +++++++ tools/tests/fake_nova_data_cli.py | 140 +++++++++ tools/tests/lib.sh | 59 ++++ tools/tests/run_all.sh | 46 +++ tools/tests/scenario_concurrency.sh | 69 +++++ tools/tests/scenario_crash_restart.sh | 88 ++++++ tools/tests/scenario_local_backlog.sh | 94 ++++++ tools/tests/scenario_local_live.sh | 82 ++++++ tools/tests/scenario_real_smoke.sh | 84 ++++++ tools/tests/scenario_remote_backlog.sh | 68 +++++ tools/tests/scenario_remote_live.sh | 87 ++++++ tools/tests/scenario_sizing.sh | 50 ++++ tools/validate_batch.py | 57 ++++ 19 files changed, 1718 insertions(+), 194 deletions(-) create mode 100644 tools/AGENT.md create mode 100755 tools/pipeline.sh delete mode 100755 tools/sync_loop.sh create mode 100644 tools/tests/README.md create mode 100644 tools/tests/fake_nova_data_cli.py create mode 100755 tools/tests/lib.sh create mode 100755 tools/tests/run_all.sh create mode 100755 tools/tests/scenario_concurrency.sh create mode 100755 tools/tests/scenario_crash_restart.sh create mode 100755 tools/tests/scenario_local_backlog.sh create mode 100755 tools/tests/scenario_local_live.sh create mode 100755 tools/tests/scenario_real_smoke.sh create mode 100755 tools/tests/scenario_remote_backlog.sh create mode 100755 tools/tests/scenario_remote_live.sh create mode 100755 tools/tests/scenario_sizing.sh create mode 100755 tools/validate_batch.py diff --git a/README.md b/README.md index 715b2d7..59b3f13 100644 --- a/README.md +++ b/README.md @@ -43,30 +43,38 @@ cameras. Examples: [`examples/`](examples/). Schema: config field, the export formats, camera resizing, and how the trimming modes choose episode boundaries (with figures). -## Live sync + export (`tools/sync_loop.sh`) - -Pulls recordings from a remote machine while collection is still running there, -and exports them in the background as each one completes — instead of waiting -for collection to finish before exporting anything. - -- **Two machines (collector + this one):** requires passwordless SSH to the - remote host (`ssh-copy-id`), since the script polls it every `POLL_SECONDS` - via `rsync`/`ssh`. Edit `REMOTE_HOST`, `REMOTE_DIRS`, and `LOCAL_DEST` at the - top of the script first. -- **Same machine:** SSH isn't needed if collection and export run on one box — - point `REMOTE_DIRS`/`LOCAL_DEST` at local paths and swap `sync_once`'s - `rsync` for a local copy (or skip syncing and export straight from the - collection dir). Not built yet; the script currently assumes a remote host. -- The script stops polling once the remote dir has been idle for - `IDLE_MINUTES`, does one final sync + export pass, then merges all batches - into one LeRobot dataset via `tools/merge_batches.py` (nova-data-cli itself - has no incremental/append mode, so each batch is a separate `--output` dir - until merged). +## Live sync + parallel export (`tools/pipeline.sh`) + +Pulls recordings while collection is still running and exports them in +parallel as they finish, instead of waiting for collection to end and then +exporting one at a time. Worker count and memory budget are computed from the +machine's own RAM/cores at startup. See [`tools/AGENT.md`](tools/AGENT.md) for +the full design (crash-safety, bisection, why it's shaped this way). + +**Two machines** (collector elsewhere, export runs here) — needs passwordless +SSH to the collector (`ssh-copy-id`): + +```bash +tools/pipeline.sh # default: --mode remote +``` + +Edit `REMOTE_HOST`/`REMOTE_DIRS`/`WATCH_DIR` at the top of the script, or +override per-run via `PIPELINE_REMOTE_HOST`, `PIPELINE_REMOTE_DIRS`, etc. +(every config value is a `PIPELINE_*` env var — see the top of the script). + +**One machine** (collection already finished, or writing straight into a +local dir) — no network involved: ```bash -tools/sync_loop.sh +PIPELINE_MODE=local PIPELINE_WATCH_DIR=/path/to/recordings tools/pipeline.sh ``` +Both modes end the same way: once nothing new has shown up for +`PIPELINE_IDLE_MINUTES` (or immediately, for a backlog that was never live), +it merges every batch into one dataset at `_merged`. Restarting +after a crash/kill is always safe — already-exported recordings are never +redone. + ## Tests ```bash diff --git a/src/nova_export/cli.py b/src/nova_export/cli.py index a3df58b..6ab2162 100644 --- a/src/nova_export/cli.py +++ b/src/nova_export/cli.py @@ -26,6 +26,7 @@ from __future__ import annotations import argparse +import json import os import shutil import subprocess @@ -204,6 +205,24 @@ def on_progress(current: int, total: int) -> None: result.output_dir, ) + # Machine-readable per-invocation outcome (keyed by segment, not recording; + # see tools/AGENT.md for how callers work around that). + metadata = result.metadata or {} + (args.output / "export_summary.json").write_text( + json.dumps( + { + "total_episodes_attempted": metadata.get("total_episodes_attempted", 0), + "successful_episodes": metadata.get("successful_episodes", 0), + "skipped_episodes": metadata.get("skipped_episodes", 0), + "failed_episodes": metadata.get("failed_episodes", 0), + "successful_list": metadata.get("successful_list", []), + "skipped_list": metadata.get("skipped_list", []), + "failed_list": metadata.get("failed_list", []), + }, + indent=2, + ) + ) + if config.format == "groot": _convert_to_groot_v21(args.output) diff --git a/tools/AGENT.md b/tools/AGENT.md new file mode 100644 index 0000000..0dfe011 --- /dev/null +++ b/tools/AGENT.md @@ -0,0 +1,176 @@ +# `pipeline.sh` design notes + +This documents *why* `tools/pipeline.sh` (+ `validate_batch.py`, `merge_batches.py`) +is built the way it is. For usage, see the main [README](../README.md). This +file is for whoever next has to change this script and needs the reasoning, +not just the code. + +## The problem + +`nova-data-cli` has no incremental/append mode — every `--output` must be a +fresh directory — and a single invocation is single-threaded and CPU-light +(~1-1.7 cores, ~50s/episode observed). Collecting ~1000 episodes and exporting +them sequentially after collection finishes wastes hours: pulling and +exporting can overlap, and multiple exports can run concurrently on idle +cores. `pipeline.sh` does both, then merges the results into one dataset. + +## Why a claim/commit protocol instead of static work partitioning + +An earlier, simpler design statically split the recording list across N +workers up front. That doesn't hold up once new recordings can arrive mid-run +(the live-feed case) — a static split would need recomputing on every arrival, +which is the same problem in a different shape. Instead, workers dynamically +claim work from a shared pool: + +- **Claiming** is `mkdir claimed/` — atomic on any POSIX filesystem + (including NFS, unlike `open(O_EXCL)`), so no separate lock is needed to + prevent two workers claiming the same recording. +- **"Done" is derived, not tracked.** There's no separate done-marker + directory. A batch's claimed IDs are written into `.claimed_ids` *inside* + its own tmp output dir *before* exporting, so the single atomic `mv` that + commits the batch to its final path is simultaneously the durable record of + every ID it contains. `rebuild_done()` just reads `.claimed_ids` out of + every existing `batch_*` dir. This collapses what would otherwise be two + separate operations (commit, then mark-done) into one atomic one — and that + matters: an earlier version *did* do them separately (commit, then loop + `touch done/`), and a crash between the two left a committed-but- + unmarked batch, which got silently re-exported and duplicated on retry. One + atomic operation can't have that window. + +## The race this still has, and how it's closed + +`rebuild_done()` is a point-in-time snapshot. If it's read at the top of a +worker's loop, then *before* that worker's `mkdir` lands, another worker +commits and releases a claim on one of the same IDs, the snapshot is stale — +the ID looks unclaimed (it is) and undone (it isn't, but the snapshot doesn't +know that yet). The `mkdir` would then succeed and re-export something already +committed. This was reproduced directly (via temporary tracing) during +development: worker A committed `rec05/06/07` and released their claims 10ms +before worker B's claim scan reached them. + +The fix is one fresh `rebuild_done()` call *after* claiming, not before: +`mv` (commit) always fully completes before `rmdir` (release) for a given ID, +so the moment a claim becomes available, any commit for that ID is already on +disk. And once a worker holds a claim, nobody else can commit that ID (commit +requires a claim to run at all) — so a single recheck immediately after +claiming can never go stale again. This is why the recheck is a *quality* +distinct from "check more often": one recheck at the right point is +permanently sufficient, not just less likely to race. + +## Bisection instead of batch-wide quarantine + +`nova-data-cli` exits 0 even when it silently skips or fails individual +episodes within a batch — exit code alone doesn't mean "all N claimed +recordings are in the output." `validate_batch.py` checks +`export_summary.json` against `.claimed_ids`, but that file is keyed by +*segment*, not recording (a single `.rrd` can yield multiple segments, and the +exporter has no per-segment source tag) — so a skip/fail inside a multi-ID +batch can't be attributed to a specific recording from that file alone. + +Rather than quarantine (or blindly retry) the whole batch — which would +needlessly punish healthy batch-mates for one bad recording — a batch that +fails validation is bisected: split in half, each half retried independently, +recursing until it narrows to size 1. At size 1, a skip/fail is unambiguous +("the only recording in this invocation didn't produce an episode"), so the +per-ID failure counter and 3-strikes quarantine only apply there. Counting +failures at every bisection level would let one real failure rack up several +"strikes" on the way down and quarantine a perfectly good recording that +happened to share unlucky batches. + +The failure counter itself is a single-byte atomic append +(`printf x >> failed/`, count = file size) recorded *at claim time*, not +after a failure — so a hard kill/OOM of the worker mid-export still counts as +an attempt, rather than letting a recording that keeps crashing the worker +retry forever without ever reaching the quarantine threshold. + +## Crash safety in general + +- **Single-instance lock**: the supervisor holds an `flock` for its entire + life. Because of that, on startup, any leftover `claimed/`/`.tmp/` entries + are provably abandoned (nothing else can legitimately be running) — except + a killed run's process group might still have orphaned children alive even + though the lock is free again; that's checked separately via a recorded + pgid before sweeping. +- **Process group teardown**: the supervisor `setsid`s itself once at startup + (becoming its own group leader) so a single `kill -- -$$` in its shutdown + trap reaches every acquire/worker/`nova-data-cli`/ffmpeg descendant, not + just its direct children. +- **Restart is always safe**: nothing needs manual cleanup after a crash. + Committed batches stay committed (derived done-state), abandoned claims are + swept, in-flight work is simply redone (at most `CHUNK` recordings' worth + per worker, since claim is held for the whole bisection tree of one claim + round). + +## Acquisition modes + +`--mode remote|local` share every line of the claim/export/validate/commit/ +merge machinery — the only thing that differs is how new `recording.rrd` +files show up in `$WATCH_DIR` and how "collection is finished" is detected: + +- **remote**: `rsync`/`ssh` on a poll loop; idle-detected via the remote + host's own max file mtime being unchanged for `IDLE_MINUTES`. `rsync` exit + codes 23/24 ("partial transfer"/"some files vanished") are expected — the + source is being actively written — and don't abort the loop. A failed or + timed-out `ssh` idle-check is treated as *inconclusive*, not as evidence of + idleness (a network blip must never cause an early merge) and not as fatal + (a network blip must never abort the whole pipeline either). +- **local**: no network at all — idle-detected via a local `find` on + `$WATCH_DIR`'s own mtimes. Since there's no rsync-provided atomicity for a + locally-written file (rsync's `--partial` semantics meant a file only + appears at its final name once fully transferred), local mode additionally + requires `recording.rrd` to be untouched for 60s before treating it as + finished writing, as a substitute completeness signal. + +This is also why **switching a finished remote collection to `--mode local`** +before going fully offline is the correct move, not a workaround: remote +mode's "safe to merge" signal fundamentally requires reaching the remote host +to confirm nothing's still arriving, so it can never fire without network — +whereas local mode's signal is purely local file-mtime watching and works +fully offline, correctly, once nothing new can possibly arrive. + +## Resource budgeting + +Worker count and per-worker memory cap are computed at startup from +`/proc/meminfo`/`nproc`, not hardcoded, so this doesn't silently misbehave on +a different machine. `MEM_TARGET_FRACTION` bounds total worker memory as a +fraction of system RAM (default 0.55 — deliberately conservative, since the +per-worker estimate is a rough baseline, not a guarantee, and the fraction is +computed against *total* memory, not what's actually free right now). Each +worker's `nova-data-cli` runs inside a `systemd-run --scope` cgroup with a +hard `MemoryMax` and `MemorySwapMax=0` — this is what actually prevents the +failure mode observed during development, where a single oversized batch +caused an OOM-kill of an internal subprocess (`rerun`) that left a truncated, +silently-corrupt output directory. A cgroup-enforced kill is a clean, cheap, +retryable failure; letting the kernel's OOM-killer pick an arbitrary victim +under memory pressure is not. A live watchdog additionally pauses claiming +new work when `MemAvailable` drops near what one worker needs, independent of +the startup sizing — this reacts to whatever else is running on the machine +right now, not just what the pipeline itself is doing. + +`nice -n 10`/`ionice -c2 -n7` on the export command is the mechanism that +keeps this from monopolizing CPU/disk even without an active watchdog for +those resources — it tells the kernel to prefer any other process, so the +pipeline only consumes spare capacity. + +## Merge + +`merge_batches.py` merges via `lerobot.datasets.aggregate.aggregate_datasets`, +which has two gaps this script covers: + +1. Its own metadata check (`validate_all_metadata`) compares fps/robot_type/ + features, but `features_equal_for_merge` *strips* video-encoder info + (codec/pix_fmt/resolution) before that comparison — so encoder drift + between batches passes validation and only surfaces deep inside the + expensive video-concatenation copy, potentially hours in. This script + checks those fields itself, up front, before starting. +2. It has no atomic-output guarantee — a killed merge leaves a partial, + multi-GB output directory that can't just be resumed (its own append logic + assumes a from-scratch destination). `merge_batches.py` builds into a + `.tmp-` directory and `os.rename`s it into place only on + success, matching the same commit pattern used for individual batches. + +The merge itself runs exactly once, after the acquisition process and every +worker have drained (checked by relaunching the worker pool if a pass finds +leftover candidates — a straggler bisection retry can land right as the last +worker exits — until a pass finds nothing left; there's no cap on this, since +remaining work only ever shrinks). diff --git a/tools/merge_batches.py b/tools/merge_batches.py index abbd770..1a794de 100644 --- a/tools/merge_batches.py +++ b/tools/merge_batches.py @@ -1,36 +1,103 @@ #!/usr/bin/env python -"""Merge the batch_* LeRobot dataset dirs produced by sync_loop.sh into one dataset.""" +"""Merge the batch_*/dataset dirs produced by tools/pipeline.sh into one dataset. + +See AGENT.md for why this checks video-encoder compatibility itself rather +than relying on lerobot's aggregate_datasets (which ignores it). +""" import argparse +import json +import os +import shutil import sys from pathlib import Path from lerobot.datasets.aggregate import aggregate_datasets +_VIDEO_INFO_KEYS_TO_CHECK = ("video.codec", "video.pix_format", "video.height", "video.width") + + +def _dataset_dirs(batches_root: Path) -> list[Path]: + return sorted(p / "dataset" for p in batches_root.glob("batch_*") if (p / "dataset").is_dir()) + + +def _load_info(dataset_dir: Path) -> dict: + return json.loads((dataset_dir / "meta" / "info.json").read_text()) + + +def _check_compatible(dataset_dirs: list[Path]) -> None: + """Fail fast on fps/robot_type/video-encoder drift.""" + first_info = _load_info(dataset_dirs[0]) + first_fps = first_info.get("fps") + first_robot_type = first_info.get("robot_type") + first_features = first_info.get("features", {}) + + problems = [] + for d in dataset_dirs[1:]: + info = _load_info(d) + if info.get("fps") != first_fps: + problems.append(f"{d}: fps={info.get('fps')} != {first_fps}") + if info.get("robot_type") != first_robot_type: + problems.append(f"{d}: robot_type={info.get('robot_type')} != {first_robot_type}") + + features = info.get("features", {}) + for key, feat in first_features.items(): + if feat.get("dtype") != "video": + continue + other_feat = features.get(key, {}) + first_video_info = feat.get("info") or {} + other_video_info = other_feat.get("info") or {} + for video_key in _VIDEO_INFO_KEYS_TO_CHECK: + if first_video_info.get(video_key) != other_video_info.get(video_key): + problems.append( + f"{d}: feature '{key}' {video_key}=" + f"{other_video_info.get(video_key)!r} != {first_video_info.get(video_key)!r}" + ) + + if problems: + print( + "Refusing to merge — batches disagree on schema/video encoding:\n " + + "\n ".join(problems), + file=sys.stderr, + ) + sys.exit(1) + def main(): parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--batches-root", required=True, type=Path, help="Dir containing batch_* dataset dirs") + parser.add_argument("--batches-root", required=True, type=Path, help="Dir containing batch_* dirs") parser.add_argument("--output", required=True, type=Path, help="Output dir for the merged dataset") parser.add_argument("--repo-id", default="merged", help="Identifier for the merged dataset") args = parser.parse_args() - batches = sorted(p for p in args.batches_root.glob("batch_*") if p.is_dir()) - if not batches: - print(f"No batch_* dirs found under {args.batches_root}", file=sys.stderr) + dataset_dirs = _dataset_dirs(args.batches_root) + if not dataset_dirs: + print(f"No batch_*/dataset dirs found under {args.batches_root}", file=sys.stderr) sys.exit(1) if args.output.exists(): print(f"Output dir already exists, refusing to overwrite: {args.output}", file=sys.stderr) sys.exit(1) - print(f"Merging {len(batches)} batches into {args.output}") - aggregate_datasets( - repo_ids=[b.name for b in batches], - aggr_repo_id=args.repo_id, - roots=batches, - aggr_root=args.output, - ) + _check_compatible(dataset_dirs) + + tmp_output = args.output.parent / f"{args.output.name}.tmp-{os.getpid()}" + if tmp_output.exists(): + shutil.rmtree(tmp_output) + + print(f"Merging {len(dataset_dirs)} batches into {args.output} (via {tmp_output})") + try: + aggregate_datasets( + repo_ids=[d.parent.name for d in dataset_dirs], + aggr_repo_id=args.repo_id, + roots=dataset_dirs, + aggr_root=tmp_output, + ) + except BaseException: + shutil.rmtree(tmp_output, ignore_errors=True) + raise + + os.rename(tmp_output, args.output) print("Merge complete.") diff --git a/tools/pipeline.sh b/tools/pipeline.sh new file mode 100755 index 0000000..afdfdb5 --- /dev/null +++ b/tools/pipeline.sh @@ -0,0 +1,392 @@ +#!/usr/bin/env bash +# Pulls episode recordings (remote over rsync/ssh, or a local dir) and exports +# them in parallel as they finish, instead of waiting for collection to end. +# See README.md for usage, AGENT.md for how/why this is built the way it is. +set -euo pipefail +shopt -s nullglob + +# ---- config ----------------------------------------------------------- +# Every value is overridable via a PIPELINE_* env var; see README.md. +MODE="${PIPELINE_MODE:-remote}" +REMOTE_HOST="${PIPELINE_REMOTE_HOST:-intern@172.31.11.129}" +if [[ -n "${PIPELINE_REMOTE_DIRS:-}" ]]; then + IFS=':' read -r -a REMOTE_DIRS <<< "$PIPELINE_REMOTE_DIRS" +else + REMOTE_DIRS=( + "/mnt/data/sebastian/raw_datasets/pick_and_place_sim_20260805_191245" + ) +fi +WATCH_DIR="${PIPELINE_WATCH_DIR:-/home/sebi/ws/Data/raw_data/choreo2/pick_and_place_sim_20260805_191245}" + +NOVA_CLI_DIR="${PIPELINE_NOVA_CLI_DIR:-/home/sebi/ws/nova-data-cli}" +EXPORT_CONFIG="${PIPELINE_EXPORT_CONFIG:-/home/sebi/ws/pick_and_place_imitation_learning/data_collection/configs/lerobot_export.json}" +EXPORT_ROOT="${PIPELINE_EXPORT_ROOT:-/home/sebi/ws/Data/choreo2_export}" +read -r -a EXPORT_CLI_CMD <<< "${PIPELINE_EXPORT_CLI:-uv run nova-data-cli}" # swap in a stub for tests + +CHUNK="${PIPELINE_CHUNK:-8}" +POLL_SECONDS="${PIPELINE_POLL_SECONDS:-180}" +IDLE_MINUTES="${PIPELINE_IDLE_MINUTES:-10}" + +# Worker count/memory cap are computed at startup from this machine's actual +# resources, not hardcoded — tune these two, not compute_workers() below. +MEM_TARGET_FRACTION="${PIPELINE_MEM_TARGET_FRACTION:-0.55}" +WORKER_MEM_ESTIMATE_MB="${PIPELINE_WORKER_MEM_ESTIMATE_MB:-3200}" + +STATE="${EXPORT_ROOT}/.pipeline" +LOCK="${STATE}/lock" +CLAIMED="${STATE}/claimed" +FAILED="${STATE}/failed" +QUARANTINE="${STATE}/quarantine" +SCRATCH="${STATE}/scratch" +TMP_OUT="${EXPORT_ROOT}/.tmp" +LOGS="${STATE}/logs" +COLLECTION_DONE="${STATE}/collection_done" +PGID_FILE="${STATE}/supervisor.pgid" + +# ---- arg parsing -------------------------------------------------------- +ROLE="${1:-supervisor}" +shift || true +POSITIONAL=() +while [[ $# -gt 0 ]]; do + case "$1" in + --mode) MODE="$2"; shift 2 ;; + *) POSITIONAL+=("$1"); shift ;; + esac +done +set -- "${POSITIONAL[@]}" + +if [[ "$MODE" == "remote" ]]; then + [[ ${#REMOTE_DIRS[@]} -eq 0 ]] && { echo "remote mode needs REMOTE_DIRS" >&2; exit 1; } + WATCH_DIR="$(dirname "$WATCH_DIR")/$(basename "${REMOTE_DIRS[0]}")" +elif [[ "$MODE" == "local" ]]; then + [[ -d "$WATCH_DIR" ]] || { echo "local mode needs WATCH_DIR to already exist: $WATCH_DIR" >&2; exit 1; } +else + echo "Unknown --mode: $MODE (expected remote|local)" >&2; exit 1 +fi + +mkdir -p "$STATE" "$CLAIMED" "$FAILED" "$QUARANTINE" "$SCRATCH" "$TMP_OUT" "$LOGS" + +# ---- resource sizing (supervisor computes once, workers inherit via env) -- +compute_workers() { + local mem_total_kb mem_workers core_workers + mem_total_kb="$(awk '/MemTotal/{print $2}' /proc/meminfo)" + mem_workers="$(awk -v kb="$mem_total_kb" -v f="$MEM_TARGET_FRACTION" -v est="$WORKER_MEM_ESTIMATE_MB" \ + 'BEGIN{printf "%d", (kb/1024*f)/est}')" # awk not $(( )): bash treats leading-zero numbers as octal + core_workers=$(( $(nproc) * 8 / 10 )) # leave headroom, don't claim every core + local workers=$mem_workers + [[ $core_workers -lt $workers ]] && workers=$core_workers + [[ $workers -lt 1 ]] && workers=1 + echo "$workers" +} + +mem_available_floor_mb() { + awk -v est="$WORKER_MEM_ESTIMATE_MB" 'BEGIN{printf "%d", est * 1.2}' +} + +# ---- shared helpers ------------------------------------------------------ +log() { echo "[$(date +%H:%M:%S)] [$ROLE] $*"; } + +mem_available_mb() { + awk '/MemAvailable/{printf "%d", $2/1024}' /proc/meminfo +} + +wait_for_memory() { + local floor; floor="$(mem_available_floor_mb)" + while [[ "$(mem_available_mb)" -lt "$floor" ]]; do + log "MemAvailable below ${floor}MB, waiting for headroom before claiming more work" + sleep 15 + done +} + +# "done" is derived from .claimed_ids in every committed batch, not a marker +# file — see AGENT.md for why. +declare -A DONE_IDS +rebuild_done() { + DONE_IDS=() + local batch f id + for batch in "$EXPORT_ROOT"/batch_*; do + f="$batch/.claimed_ids" + [[ -f "$f" ]] || continue + while IFS= read -r id; do + [[ -n "$id" ]] && DONE_IDS["$id"]=1 + done < "$f" + done +} + +is_candidate() { + local id="$1" rrd="$WATCH_DIR/$id/recording.rrd" + [[ -f "$rrd" ]] || return 1 + [[ -n "${DONE_IDS[$id]:-}" ]] && return 1 + [[ -d "$CLAIMED/$id" ]] && return 1 + [[ -f "$QUARANTINE/$id" ]] && return 1 + if [[ "$MODE" == "local" ]]; then + [[ -n "$(find "$rrd" -mmin +1 2>/dev/null)" ]] || return 1 # must be untouched 60s (no rsync atomicity here) + fi + return 0 +} + +list_candidates() { + local id + for d in "$WATCH_DIR"/*/; do + id="$(basename "$d")" + is_candidate "$id" && echo "$id" + done +} + +failure_count() { local f="$FAILED/$1"; [[ -f "$f" ]] && wc -c < "$f" || echo 0; } +record_attempt() { printf x >> "$FAILED/$1"; } + +# ---- acquisition --------------------------------------------------------- +role_acquire() { + mkdir -p "$WATCH_DIR" + if [[ "$MODE" == "local" ]]; then + log "local mode: watching $WATCH_DIR, no network step" + local idle_polls_needed=$(( (IDLE_MINUTES * 60 + POLL_SECONDS - 1) / POLL_SECONDS )) idle_count=0 + while true; do + if [[ -z "$(find "$WATCH_DIR" -type f -mmin "-${IDLE_MINUTES}" -print -quit 2>/dev/null)" ]]; then + idle_count=$((idle_count + 1)) + [[ $idle_count -ge $idle_polls_needed ]] && break + else + idle_count=0 + fi + sleep "$POLL_SECONDS" + done + touch "$COLLECTION_DONE" + log "local collection idle for ${IDLE_MINUTES}m, done" + return + fi + + local ssh_opts=(-o BatchMode=yes -o ConnectTimeout=10 -o ServerAliveInterval=15 -o ServerAliveCountMax=4) + local idle_polls_needed=$(( (IDLE_MINUTES * 60 + POLL_SECONDS - 1) / POLL_SECONDS )) idle_count=0 last_mtime="" + + sync_once() { + local rc=0 + rsync -a --partial --info=stats2 --timeout=300 \ + "${REMOTE_DIRS[@]/#/$REMOTE_HOST:}" "$(dirname "$WATCH_DIR")/" || rc=$? + case "$rc" in + 0|23|24) ;; # 23/24: source still being written, expected + *) log "rsync exited $rc (non-fatal, will retry next poll)" ;; + esac + } + + remote_max_mtime() { + timeout 60 ssh "${ssh_opts[@]}" "$REMOTE_HOST" \ + "find ${REMOTE_DIRS[*]} -type f -printf '%T@\n' 2>/dev/null | sort -n | tail -1" + } + + while true; do + sync_once + local mtime rc=0 + mtime="$(remote_max_mtime)" || rc=$? + if [[ $rc -ne 0 || -z "$mtime" ]]; then + log "idle-check ssh/find failed or timed out — inconclusive, not counted toward idle" + elif [[ "$mtime" == "$last_mtime" ]]; then + idle_count=$((idle_count + 1)) + else + idle_count=0 + last_mtime="$mtime" + fi + + if [[ $idle_count -ge $idle_polls_needed ]]; then + log "no new remote data for ${IDLE_MINUTES}m, final sync" + sync_once + break + fi + sleep "$POLL_SECONDS" + done + touch "$COLLECTION_DONE" +} + +# ---- worker --------------------------------------------------------------- +# Claims up to CHUNK candidates, exports them in one nova-data-cli invocation, +# validates, then atomically commits. See AGENT.md for the full protocol. +export_and_commit() { + local -a ids=("$@") + local batch_name="batch_$(date +%Y%m%d_%H%M%S)_w${WORKER_IDX}_$$_${RANDOM}" + local scratch="$SCRATCH/$batch_name" out="$TMP_OUT/$batch_name" + + rm -rf "$scratch" + mkdir -p "$scratch" "$out" + printf '%s\n' "${ids[@]}" > "$out/.claimed_ids" + local id + for id in "${ids[@]}"; do + ln -s "$WATCH_DIR/$id" "$scratch/$id" + done + + local ok=1 + # nice/ionice: yield cores/disk to other processes instead of a CPU watchdog. + # PYTHONUNBUFFERED: so progress lines stream live instead of buffering. + local export_cmd=(env PYTHONUNBUFFERED=1 nice -n 10 ionice -c2 -n7 "${EXPORT_CLI_CMD[@]}" --dataset "$scratch" --config "$EXPORT_CONFIG" --output "$out/dataset") + if command -v systemd-run >/dev/null 2>&1; then + (cd "$NOVA_CLI_DIR" && systemd-run --user --scope -p "MemoryMax=${WORKER_MEM_ESTIMATE_MB}M" -p MemorySwapMax=0 --collect -- "${export_cmd[@]}") || ok=0 + else + log "systemd-run not available — running without a memory cgroup cap (best effort only)" + (cd "$NOVA_CLI_DIR" && "${export_cmd[@]}") || ok=0 + fi + + if [[ $ok -eq 1 ]]; then + if uv run python "$NOVA_CLI_DIR/tools/validate_batch.py" --output-dir "$out/dataset" --claimed-ids "$out/.claimed_ids"; then + mv -T "$out" "$EXPORT_ROOT/$batch_name" + rm -rf "$scratch" + log "committed $batch_name (${#ids[@]} recording(s))" + return 0 + fi + fi + + rm -rf "$out" "$scratch" + return 1 +} + +# Failure counter/quarantine only apply at batch size 1 (see AGENT.md for why). +attempt_batch() { + local -a ids=("$@") + + if [[ ${#ids[@]} -eq 1 ]]; then + local id="${ids[0]}" + record_attempt "$id" # before running, so a hard kill/OOM still counts + if export_and_commit "$id"; then + rmdir "$CLAIMED/$id" 2>/dev/null || true + else + if [[ "$(failure_count "$id")" -ge 3 ]]; then + log "quarantining $id after 3 failed attempts" + touch "$QUARANTINE/$id" + rm -f "$FAILED/$id" + fi + rmdir "$CLAIMED/$id" 2>/dev/null || true + fi + return + fi + + if export_and_commit "${ids[@]}"; then + local id + for id in "${ids[@]}"; do rmdir "$CLAIMED/$id" 2>/dev/null || true; done + return + fi + + local half=$(( ${#ids[@]} / 2 )) + local -a left=("${ids[@]:0:half}") right=("${ids[@]:half}") + log "batch of ${#ids[@]} failed/ambiguous, bisecting into ${#left[@]} + ${#right[@]}" + attempt_batch "${left[@]}" + attempt_batch "${right[@]}" +} + +role_worker() { + WORKER_IDX="$1" + while true; do + rebuild_done + wait_for_memory + + local -a claimed=() + local id + while IFS= read -r id; do + [[ ${#claimed[@]} -ge $CHUNK ]] && break + mkdir "$CLAIMED/$id" 2>/dev/null && claimed+=("$id") + done < <(list_candidates) + + # Recheck against a fresh rebuild_done: another worker may have committed + # (and released) one of these IDs in the gap since our scan (see AGENT.md). + if [[ ${#claimed[@]} -gt 0 ]]; then + rebuild_done + local -a fresh=() + for id in "${claimed[@]}"; do + if [[ -n "${DONE_IDS[$id]:-}" ]]; then + rmdir "$CLAIMED/$id" 2>/dev/null || true + else + fresh+=("$id") + fi + done + claimed=("${fresh[@]}") + fi + + if [[ ${#claimed[@]} -eq 0 ]]; then + if [[ -f "$COLLECTION_DONE" ]]; then + log "no candidates and collection done, exiting" + return 0 + fi + sleep 30 + continue + fi + + attempt_batch "${claimed[@]}" + done +} + +# ---- supervisor ----------------------------------------------------------- +role_supervisor() { + # setsid makes this the leader of a fresh process group so `kill -- -$$` on + # shutdown reaches every descendant (see AGENT.md). + if [[ -z "${PIPELINE_RESPAWNED:-}" ]]; then + exec env PIPELINE_RESPAWNED=1 setsid "$0" supervisor --mode "$MODE" + fi + + exec 9>"$LOCK" + if ! flock -n 9; then + echo "Another pipeline.sh is already running (lock held: $LOCK)" >&2 + exit 1 + fi + + # Holding the lock means it's safe to sweep stale state, unless a killed + # run's process group is somehow still alive. + if [[ -f "$PGID_FILE" ]]; then + old_pgid="$(cat "$PGID_FILE")" + if [[ -n "$old_pgid" ]] && kill -0 -- "-$old_pgid" 2>/dev/null; then + echo "Previous run's process group ($old_pgid) still alive — kill it before restarting" >&2 + exit 1 + fi + fi + rm -rf "${CLAIMED:?}"/* "${TMP_OUT:?}"/* "${SCRATCH:?}"/* + rm -f "$COLLECTION_DONE" + mkdir -p "$CLAIMED" + + echo "$$" > "$PGID_FILE" # setsid above made pid == pgid + + WORKERS="$(compute_workers)" + log "sizing: $(nproc) cores, $(awk '/MemTotal/{printf "%.1fGB", $2/1024/1024}' /proc/meminfo) RAM, MEM_TARGET_FRACTION=$MEM_TARGET_FRACTION, WORKER_MEM_ESTIMATE_MB=$WORKER_MEM_ESTIMATE_MB -> WORKERS=$WORKERS" + + trap 'log "shutting down"; kill -- -$$ 2>/dev/null || true' INT TERM + + # Tee each role's output to both terminal (prefixed) and its log file. + ( "$0" acquire --mode "$MODE" 2>&1 | sed -u 's/^/[acquire] /' | tee -a "$LOGS/acquire.log" ) & + local acquire_pid=$! + + local -a worker_pids=() + spawn_workers() { + worker_pids=() + local i + for ((i = 0; i < WORKERS; i++)); do + ( "$0" worker --mode "$MODE" "$i" 2>&1 | sed -u "s/^/[w${i}] /" | tee -a "$LOGS/w${i}.log" ) & + worker_pids+=($!) + done + } + spawn_workers + wait "$acquire_pid" || true + wait "${worker_pids[@]}" || true + + # Drain: relaunch until a pass finds nothing left (see AGENT.md). + while true; do + rebuild_done + if [[ -z "$(list_candidates | head -1)" ]] && [[ -z "$(ls -A "$CLAIMED" 2>/dev/null)" ]]; then + break + fi + log "drain pass found leftover work, relaunching workers" + spawn_workers + wait "${worker_pids[@]}" || true + done + + local -a batch_dirs=("$EXPORT_ROOT"/batch_*) + [[ ${#batch_dirs[@]} -gt 0 ]] || { log "no batches produced, nothing to merge"; return 0; } + + log "all workers drained, merging" + (cd "$NOVA_CLI_DIR" && uv run python tools/merge_batches.py \ + --batches-root "$EXPORT_ROOT" \ + --output "${EXPORT_ROOT}_merged") + log "done: ${EXPORT_ROOT}_merged" +} + +case "$ROLE" in + supervisor) role_supervisor ;; + acquire) role_acquire ;; + worker) role_worker "${1:?worker index required}" ;; + sizing) compute_workers ;; # test hook: print the computed WORKERS count and exit + *) echo "Unknown role: $ROLE" >&2; exit 1 ;; +esac diff --git a/tools/sync_loop.sh b/tools/sync_loop.sh deleted file mode 100755 index 9b03da9..0000000 --- a/tools/sync_loop.sh +++ /dev/null @@ -1,162 +0,0 @@ -#!/usr/bin/env bash -# Pull recordings from the WS while collection is still running there, and -# export them in the background as they arrive, so pulling + exporting happen -# in parallel instead of waiting for all 1000 episodes then exporting. -# -# A recording is only "done" once /recording.rrd exists (it's -# written once on stop; chunks/ alone means still recording) — that's the -# completeness check, both for what's safe to export. -# -# nova-data-cli has no incremental/append mode (--output must not exist), so -# this can't grow one dataset live: each newly-arrived batch of recordings -# gets exported into its own batch_N output dir. Once collection is finished -# (idle timeout below), the batches are merged into one final LeRobot dataset -# via lerobot's own aggregate_datasets (tools/merge_batches.py). -set -euo pipefail -shopt -s nullglob - -REMOTE_HOST="intern@172.31.11.129" -REMOTE_DIRS=( - "/mnt/data/sebastian/raw_datasets/pick_and_place_sim_20260805_191245" -) -LOCAL_DEST="/home/sebi/ws/Data/raw_data/choreo2" - -NOVA_CLI_DIR="/home/sebi/ws/nova-data-cli" -EXPORT_CONFIG="/home/sebi/ws/pick_and_place_imitation_learning/data_collection/configs/lerobot_export.json" -EXPORT_OUTPUT_ROOT="/home/sebi/ws/Data/choreo2_export" - -LEDGER="${LOCAL_DEST}/.exported_ids" -INFLIGHT="${LOCAL_DEST}/.inflight_ids" -EXPORT_PID_FILE="${LOCAL_DEST}/.export.pid" - -POLL_SECONDS=180 -IDLE_MINUTES=10 - -mkdir -p "$LOCAL_DEST" "$EXPORT_OUTPUT_ROOT" -touch "$LEDGER" "$INFLIGHT" - -idle_polls_needed=$(( (IDLE_MINUTES * 60 + POLL_SECONDS - 1) / POLL_SECONDS )) -idle_count=0 -last_mtime="" -batch_num=0 - -remote_max_mtime() { - ssh "$REMOTE_HOST" "find ${REMOTE_DIRS[*]} -type f -printf '%T@\n' 2>/dev/null | sort -n | tail -1" -} - -sync_once() { - rsync -avP "${REMOTE_DIRS[@]/#/$REMOTE_HOST:}" "$LOCAL_DEST/" -} - -export_running() { - [[ -f "$EXPORT_PID_FILE" ]] && kill -0 "$(cat "$EXPORT_PID_FILE")" 2>/dev/null -} - -# Recordings that have a recording.rrd (complete) and aren't already -# exported or queued in a currently-running batch. -find_new_recordings() { - for remote_dir in "${REMOTE_DIRS[@]}"; do - local_root="${LOCAL_DEST}/$(basename "$remote_dir")" - for rrd in "$local_root"/*/recording.rrd; do - recording_dir="$(dirname "$rrd")" - recording_id="$(basename "$recording_dir")" - if ! grep -qxF "$recording_id" "$LEDGER" && ! grep -qxF "$recording_id" "$INFLIGHT"; then - echo "$recording_dir" - fi - done - done -} - -run_export_batch() { - local batch_dir="$1" - shift - local recordings=("$@") - - local merged_dir="${EXPORT_OUTPUT_ROOT}/_merged_${batch_dir}" - local output_dir="${EXPORT_OUTPUT_ROOT}/${batch_dir}" - - rm -rf "$merged_dir" - mkdir -p "$merged_dir" - for recording in "${recordings[@]}"; do - ln -s "$recording" "${merged_dir}/$(basename "$recording")" - done - - if (cd "$NOVA_CLI_DIR" && uv run nova-data-cli \ - --dataset "$merged_dir" \ - --config "$EXPORT_CONFIG" \ - --output "$output_dir"); then - for recording in "${recordings[@]}"; do - basename "$recording" >> "$LEDGER" - done - echo "Exported ${#recordings[@]} recordings -> $output_dir" - else - echo "Export batch $batch_dir failed, will retry these next round: ${recordings[*]}" >&2 - fi - - for recording in "${recordings[@]}"; do - sed -i "\|^$(basename "$recording")\$|d" "$INFLIGHT" - done - rm -rf "$merged_dir" -} - -maybe_start_export_batch() { - if export_running; then - return - fi - - new_recordings=() - while IFS= read -r line; do - [[ -n "$line" ]] && new_recordings+=("$line") - done < <(find_new_recordings) - - if [[ ${#new_recordings[@]} -eq 0 ]]; then - return - fi - - batch_num=$((batch_num + 1)) - batch_name="batch_$(printf '%04d' "$batch_num")" - for recording in "${new_recordings[@]}"; do - basename "$recording" >> "$INFLIGHT" - done - - echo "Starting export batch $batch_name with ${#new_recordings[@]} new recordings" - run_export_batch "$batch_name" "${new_recordings[@]}" & - echo $! > "$EXPORT_PID_FILE" -} - -while true; do - sync_once - maybe_start_export_batch - - mtime="$(remote_max_mtime)" - if [[ "$mtime" == "$last_mtime" ]]; then - idle_count=$((idle_count + 1)) - else - idle_count=0 - last_mtime="$mtime" - fi - - if [[ $idle_count -ge $idle_polls_needed ]]; then - echo "No new remote data for ${IDLE_MINUTES}m, assuming collection finished." - break - fi - - sleep "$POLL_SECONDS" -done - -echo "Final sync + export of any remaining recordings..." -sync_once -while export_running; do - sleep 5 -done -maybe_start_export_batch -while export_running; do - sleep 5 -done - -MERGED_OUTPUT="${EXPORT_OUTPUT_ROOT}_merged" -echo "All batches exported. Merging into $MERGED_OUTPUT..." -(cd "$NOVA_CLI_DIR" && uv run python tools/merge_batches.py \ - --batches-root "$EXPORT_OUTPUT_ROOT" \ - --output "$MERGED_OUTPUT") -echo "Done. Merged dataset: $MERGED_OUTPUT" diff --git a/tools/tests/README.md b/tools/tests/README.md new file mode 100644 index 0000000..9bad96a --- /dev/null +++ b/tools/tests/README.md @@ -0,0 +1,100 @@ +# tools/pipeline.sh test suite + +Extensive tests for `tools/pipeline.sh` covering every deployment shape it +supports. None of these touch real collection data +(`/home/sebi/ws/Data/raw_data/choreo2`, `/home/sebi/ws/Data/choreo2_export`) +— every scenario runs against a disposable `/tmp/pipeline_test_*` dir (and, +for the remote scenarios, a dedicated test-only subdir under +`/mnt/data/sebastian/pipeline_test_fixtures/` on the real workstation, +never `raw_datasets/`). `lib.sh`'s `require_test_path` refuses to run any +cleanup against a path that doesn't literally contain `pipeline_test`, as a +backstop against a typo ever pointing somewhere real. + +## Running + +``` +tools/tests/run_all.sh # everything, ~15-20 min total +bash tools/tests/scenario_local_backlog.sh # any one scenario individually +``` + +Each scenario prints `PASS`/`FAIL` per assertion and a summary line; `run_all.sh` +exits nonzero if anything failed. + +## How pipeline.sh is made testable + +`pipeline.sh`'s config block reads every value from a `PIPELINE_*` env var +with the real production value as the default — setting no env vars gives +you exactly today's production behavior. Tests override `PIPELINE_MODE`, +`PIPELINE_WATCH_DIR`, `PIPELINE_EXPORT_ROOT`, `PIPELINE_REMOTE_HOST`/`_DIRS`, +`PIPELINE_CHUNK`, `PIPELINE_POLL_SECONDS`, `PIPELINE_IDLE_MINUTES`, and +`PIPELINE_EXPORT_CLI` (the nova-data-cli invocation itself — see below). +A `sizing` role was added to the role-dispatch `case` purely as a test hook: +`pipeline.sh sizing --mode local` prints the computed `WORKERS` count and exits, +letting `scenario_sizing.sh` check `compute_workers()` in isolation. + +## Two tiers + +**Tier 1 (most of the coverage): `fake_nova_data_cli.py`.** Running the real +`nova-data-cli` for every scenario/edge case would take too long (~50s/episode) +and most scenarios need to control failure/skip/timing precisely, which real +`.rrd` data can't do on demand. `PIPELINE_EXPORT_CLI` swaps it in for +`uv run nova-data-cli`. It drives the *real* `lerobot.datasets.lerobot_dataset.LeRobotDataset` +writer (with `use_videos=False` to skip ffmpeg) so its output is a genuinely +mergeable LeRobot dataset, not a hand-rolled mock of the schema — the real +`aggregate_datasets` runs for real in every scenario's merge step. + +Each fixture recording dir gets a `.behavior` file controlling what the stub +does with it: +- `success` (default) — writes one real episode +- `sleep:N` — sleeps N seconds, then succeeds (for concurrency proofs) +- `skip` — contributes 0 episodes, counted as skipped (exit 0) +- `crash` — aborts the whole batch immediately (exit 1) — mirrors the real + CLI's `SystemExit(1)` path for a config/data problem, not a per-episode failure +- `flaky:N` — skips for the first N attempts, then succeeds — proves a + transient failure recovers via retry instead of being quarantined. State is + tracked in `$FAKE_CLI_STATE_DIR` (default `/tmp/fake_nova_data_cli_state`), + deliberately *not* next to `.behavior` — that path is reached through + `pipeline.sh`'s scratch symlink back into `$WATCH_DIR`, and writing there + would reset local mode's idle-detection clock on every retry (a test-harness + artifact; the real CLI never writes back into a source recording). + +**Tier 2: `scenario_real_smoke.sh`.** One real end-to-end run through the +actual `nova-data-cli` (no stub) — two of the smallest sample recordings +under `/home/sebi/ws/Data/5-pick-cube-sim-raw/`, copied (never symlinked or +moved) into an isolated test dir, run through the real CLI, real +`export_summary.json`, real `aggregate_datasets`. Skips itself (exit 0, not a +failure) if that sample dir or the real export config isn't present on the +machine running the suite. + +## Scenarios + +| Script | Covers | +|---|---| +| `scenario_sizing.sh` | `compute_workers()` scales with `MEM_TARGET_FRACTION`/`WORKER_MEM_ESTIMATE_MB`, clamps to `[1, ~0.8*nproc]` | +| `scenario_concurrency.sh` | Genuine wall-clock proof workers overlap (`CHUNK=1`, N slow fixtures, span measured from first→last commit, compared against a fully-serial floor) | +| `scenario_local_backlog.sh` | Local mode, whole backlog present upfront ("collection already finished, still want parallelism"). Bisection isolating a bad recording without quarantining batch-mates, 3-strikes quarantine, a flaky recording recovering via retry, no duplicate/dropped recordings, correct merged episode count, >1 worker actually used | +| `scenario_local_live.sh` | Local mode, recordings trickle in via a background "collector" — export starts *before* collection finishes (explicit timestamp comparison), `collection_done` doesn't fire prematurely | +| `scenario_remote_backlog.sh` | Real SSH/rsync against the real workstation host, into a dedicated test-only remote dir, whole backlog upfront | +| `scenario_remote_live.sh` | Real SSH/rsync, recordings trickle in remotely, same overlap/idle-detection proofs as the local live scenario | +| `scenario_crash_restart.sh` | SIGKILLs the whole supervisor process group mid-run (some recordings genuinely in-flight, claimed but uncommitted), restarts against the same `EXPORT_ROOT`, verifies the restart isn't blocked by the dead process's stale pgid, stale claims get swept and requeued, and everything ends up committed exactly once (no loss, no duplicate) across both runs | +| `scenario_real_smoke.sh` | Tier 2 — real `nova-data-cli`, no stub | + +## Known timing characteristics (not bugs) + +Local-mode scenarios take ~2-4 minutes each: `is_candidate()` requires a +`recording.rrd` to sit untouched for 60s before it's claimable (no +rsync-provided atomicity locally), and idle-detection needs `IDLE_MINUTES` +(shortened to 1 in tests) of sustained quiet on top of that. Remote-mode +scenarios are faster since that 60s floor doesn't apply. `scenario_concurrency.sh` +measures only the first→last-commit span, not total wall time, specifically +to avoid that floor contaminating the timing assertion. + +## What isn't covered + +- The `systemd-run`-unavailable fallback path in `export_and_commit()` — `systemd-run` + is present on this machine, so that branch never executes here. Worth a + manual check on a machine without it. +- A crash *during* the atomic commit rename itself (vs. mid-export, which + `scenario_crash_restart.sh` does cover) isn't specifically targeted — `mv -T` + is a single rename(2) syscall, effectively instantaneous, so there's no + practical window to land a kill inside it deterministically. diff --git a/tools/tests/fake_nova_data_cli.py b/tools/tests/fake_nova_data_cli.py new file mode 100644 index 0000000..b9d1b1d --- /dev/null +++ b/tools/tests/fake_nova_data_cli.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python +"""Stand-in for `nova-data-cli`, used only by tools/tests/. + +Reads the same --dataset/--config/--output args as the real CLI. --dataset is +a scratch dir of symlinks (one per claimed recording, named by recording_id — +same layout pipeline.sh's export_and_commit() builds for the real CLI). Each +symlinked recording dir must contain a `.behavior` file controlling what this +run does with it: + + success (default if `.behavior` is missing) — writes one real episode + sleep:N — sleeps N seconds first, then behaves like `success` + skip — writes nothing for this id (counts as a skipped episode, exit 0) + crash — aborts the WHOLE invocation immediately (exit 1, no output dir + contents at all) — mirrors the real CLI's SystemExit(1) path + for a config/data problem that isn't a per-episode failure + flaky:N — behaves like `skip` for the first N invocations of this specific + recording, then `success` from then on — distinguishes + "transient, recovers on retry" from a permanently-bad recording + (`skip` forever, which should end up quarantined instead). + Attempt counts are tracked in $FAKE_CLI_STATE_DIR (default + /tmp/fake_nova_data_cli_state), NOT next to `.behavior` — that + path is reached through pipeline.sh's scratch symlink back into + $WATCH_DIR, and writing there would touch the real watched + directory's mtime on every retry, confusing local-mode's + idle-detection (a test-harness concern, not something the real + nova-data-cli would ever do — it only reads recordings). + +Uses the real lerobot.datasets.lerobot_dataset.LeRobotDataset writer (with +use_videos=False, to skip ffmpeg — an image feature is exercised just as much +of the merge path as a video one for these tests) so the dataset this writes +is genuinely mergeable by the real `lerobot.datasets.aggregate.aggregate_datasets`, +not a hand-rolled mock of the schema. + +--config's content is ignored; only its path needs to be passed through, to +match the real CLI's argv shape (pipeline.sh doesn't care what's in it). +""" + +import argparse +import json +import os +import sys +import time +from pathlib import Path + +import numpy as np + + +def read_behavior(recording_dir: Path) -> str: + f = recording_dir / ".behavior" + return f.read_text().strip() if f.is_file() else "success" + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--dataset", required=True, type=Path) + parser.add_argument("--config", required=True, type=Path) + parser.add_argument("--output", required=True, type=Path) + args = parser.parse_args() + + recording_dirs = sorted(p for p in args.dataset.iterdir() if p.is_dir()) + behaviors = {p.name: read_behavior(p) for p in recording_dirs} + + if any(b == "crash" for b in behaviors.values()): + print(f"fake_nova_data_cli: crash-tagged recording in batch {list(behaviors)}", file=sys.stderr) + return 1 + + successful, skipped, failed = [], [], [] + dataset = None + for name, behavior in behaviors.items(): + if behavior.startswith("sleep:"): + time.sleep(float(behavior.split(":", 1)[1])) + behavior = "success" + elif behavior.startswith("flaky:"): + state_dir = Path(os.environ.get("FAKE_CLI_STATE_DIR", "/tmp/fake_nova_data_cli_state")) + state_dir.mkdir(parents=True, exist_ok=True) + counter_file = state_dir / f"{name}.attempts" + attempts = int(counter_file.read_text()) + 1 if counter_file.is_file() else 1 + counter_file.write_text(str(attempts)) + threshold = int(behavior.split(":", 1)[1]) + behavior = "skip" if attempts <= threshold else "success" + + if behavior == "skip": + skipped.append(name) + continue + if behavior != "success": + failed.append(name) + continue + + if dataset is None: + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + dataset = LeRobotDataset.create( + repo_id=f"test/{args.output.name}", + fps=10, + features={ + "action": {"dtype": "float32", "shape": (2,), "names": None}, + "observation.state": {"dtype": "float32", "shape": (2,), "names": None}, + }, + root=args.output, + use_videos=False, + ) + + for _ in range(3): + dataset.add_frame( + { + "action": np.zeros(2, dtype=np.float32), + "observation.state": np.zeros(2, dtype=np.float32), + "task": "fake_task", + } + ) + dataset.save_episode() + successful.append(name) + + if dataset is None: + # Mirrors the real exporter: a batch where every segment is skipped + # raises rather than producing a valid zero-episode dataset. + print("fake_nova_data_cli: all claimed recordings skipped, nothing to export", file=sys.stderr) + return 1 + + dataset.finalize() + + (args.output / "export_summary.json").write_text( + json.dumps( + { + "total_episodes_attempted": len(behaviors), + "successful_episodes": len(successful), + "skipped_episodes": len(skipped), + "failed_episodes": len(failed), + "successful_list": successful, + "skipped_list": skipped, + "failed_list": failed, + } + ) + ) + print(f"fake_nova_data_cli: {len(successful)} ok, {len(skipped)} skipped, {len(failed)} failed") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tools/tests/lib.sh b/tools/tests/lib.sh new file mode 100755 index 0000000..df25df8 --- /dev/null +++ b/tools/tests/lib.sh @@ -0,0 +1,59 @@ +#!/usr/bin/env bash +# Shared helpers for tools/tests/scenario_*.sh. Sourced, not executed. +set -euo pipefail +shopt -s nullglob + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +PIPELINE="$REPO_ROOT/tools/pipeline.sh" +FAKE_CLI="uv run python $REPO_ROOT/tools/tests/fake_nova_data_cli.py" +REMOTE_TEST_HOST="intern@172.31.11.129" +REMOTE_TEST_ROOT="/mnt/data/sebastian/pipeline_test_fixtures" + +PASS=0 +FAIL=0 + +# Guard against ever pointing a destructive rm/ssh-rm at a real data dir — +# every test path must contain this marker. Call before any rm -rf/ssh cleanup. +require_test_path() { + case "$1" in + *pipeline_test*) ;; + *) echo "REFUSING to touch non-test-looking path: $1" >&2; exit 1 ;; + esac +} + +pass() { PASS=$((PASS + 1)); echo " PASS: $1"; } +fail() { FAIL=$((FAIL + 1)); echo " FAIL: $1"; } +assert_eq() { [[ "$1" == "$2" ]] && pass "$3 ($1)" || fail "$3 (expected [$2], got [$1])"; } + +summary() { + echo + echo "=== $PASS passed, $FAIL failed ===" + [[ $FAIL -eq 0 ]] +} + +# make_fixture +# behavior: success (default) | skip | crash | sleep:N +make_fixture() { + local kind="$1" base="$2" id="$3" behavior="${4:-success}" + if [[ "$kind" == "local" ]]; then + mkdir -p "$base/$id" + touch "$base/$id/recording.rrd" + echo "$behavior" > "$base/$id/.behavior" + else + ssh -o BatchMode=yes "$REMOTE_TEST_HOST" \ + "mkdir -p '$base/$id' && touch '$base/$id/recording.rrd' && echo '$behavior' > '$base/$id/.behavior'" + fi +} + +# All IDs actually committed into a batch anywhere under $1 (an EXPORT_ROOT). +committed_ids() { + local export_root="$1" f + for f in "$export_root"/batch_*/.claimed_ids; do + [[ -f "$f" ]] && cat "$f" + done +} + +quarantined_ids() { + local export_root="$1" + ls "$export_root/.pipeline/quarantine" 2>/dev/null || true +} diff --git a/tools/tests/run_all.sh b/tools/tests/run_all.sh new file mode 100755 index 0000000..c894bb5 --- /dev/null +++ b/tools/tests/run_all.sh @@ -0,0 +1,46 @@ +#!/usr/bin/env bash +# Runs every tools/tests/scenario_*.sh in turn and reports pass/fail per +# scenario. Each scenario is independently runnable too: bash tools/tests/scenario_X.sh +set -uo pipefail # not -e: one scenario's failure shouldn't stop the rest +cd "$(dirname "${BASH_SOURCE[0]}")" + +SCENARIOS=( + scenario_sizing.sh + scenario_concurrency.sh + scenario_local_backlog.sh + scenario_local_live.sh + scenario_remote_backlog.sh + scenario_remote_live.sh + scenario_crash_restart.sh + scenario_real_smoke.sh +) + +overall_pass=0 +overall_fail=0 +declare -A results + +for s in "${SCENARIOS[@]}"; do + echo + echo "############################################################" + echo "# $s" + echo "############################################################" + if bash "./$s"; then + results["$s"]="PASS" + overall_pass=$((overall_pass + 1)) + else + results["$s"]="FAIL" + overall_fail=$((overall_fail + 1)) + fi +done + +echo +echo "============================================================" +echo "SUMMARY" +echo "============================================================" +for s in "${SCENARIOS[@]}"; do + printf '%-30s %s\n' "$s" "${results[$s]}" +done +echo "------------------------------------------------------------" +echo "$overall_pass scenario(s) passed, $overall_fail failed" + +[[ $overall_fail -eq 0 ]] diff --git a/tools/tests/scenario_concurrency.sh b/tools/tests/scenario_concurrency.sh new file mode 100755 index 0000000..11a8a05 --- /dev/null +++ b/tools/tests/scenario_concurrency.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# Dedicated wall-clock proof that workers genuinely overlap, isolated from +# scenario_local_backlog's bisection/quarantine logic (there, batch +# composition is claim-order-dependent, so a tight timing assertion would be +# flaky). Here CHUNK=1 forces every slow fixture into its own singleton +# batch deterministically, so the concurrency proof isn't at the mercy of +# how bisection happens to group things. +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" +source ./lib.sh + +TEST_ROOT="/tmp/pipeline_test_concurrency_$$" +require_test_path "$TEST_ROOT" +trap 'kill -- "-${pid:-}" 2>/dev/null || true; rm -rf "$TEST_ROOT"' EXIT +WATCH="$TEST_ROOT/watch" +EXPORT_ROOT="$TEST_ROOT/export" +mkdir -p "$WATCH" "$EXPORT_ROOT" + +echo "=== scenario_concurrency ===" + +SLEEP_S=6 +N=6 # >= WORKERS on any reasonable machine, so every worker gets at least one +for i in $(seq -w 1 "$N"); do make_fixture local "$WATCH" "slow$i" "sleep:$SLEEP_S"; done + +export PIPELINE_MODE=local +export PIPELINE_WATCH_DIR="$WATCH" +export PIPELINE_EXPORT_ROOT="$EXPORT_ROOT" +export PIPELINE_EXPORT_CLI="$FAKE_CLI" +export PIPELINE_EXPORT_CONFIG="$TEST_ROOT/fake_config.json" +touch "$PIPELINE_EXPORT_CONFIG" +export PIPELINE_CHUNK=1 +export PIPELINE_POLL_SECONDS=3 +export PIPELINE_IDLE_MINUTES=1 + +"$PIPELINE" supervisor >"$TEST_ROOT/supervisor.log" 2>&1 & +pid=$! + +# Only wait for the N batches to commit, not for the whole idle-timeout tail +# (which would dominate wall time and isn't part of what we're proving here). +# Local mode's mandatory ~60s candidate-maturity wait (is_candidate requires +# recording.rrd to sit untouched for 60s) happens before any claim, so budget +# for that plus the actual export work. +deadline=$((SECONDS + 60 + (N * SLEEP_S) + 40)) # +40: systemd-run/claim overhead per wave, observed ~4-10s each +while [[ $(committed_ids "$EXPORT_ROOT" 2>/dev/null | wc -l) -lt $N ]]; do + [[ $SECONDS -lt $deadline ]] || break + sleep 1 +done +kill -- "-$pid" 2>/dev/null || true +wait "$pid" 2>/dev/null || true + +committed_count=$(committed_ids "$EXPORT_ROOT" 2>/dev/null | wc -l) +assert_eq "$committed_count" "$N" "all $N slow recordings committed" + +# Measure the export phase itself (first commit -> last commit), excluding +# local mode's fixed candidate-maturity wait, which is orthogonal to whether +# workers overlap. +mtimes=$(stat -c %Y "$EXPORT_ROOT"/batch_*/dataset 2>/dev/null | sort -n) +first=$(head -1 <<< "$mtimes") +last=$(tail -1 <<< "$mtimes") +span=$((last - first + 1)) # +1: same-second commits would otherwise show span=0 +serial_floor=$((N * SLEEP_S)) +echo "export phase span: ${span}s (first->last commit); fully-serial floor would be ${serial_floor}s" +if [[ $span -lt $serial_floor ]]; then + pass "export phase span (${span}s) below serial floor (${serial_floor}s) -> workers overlapped" +else + fail "export phase span (${span}s) not below serial floor (${serial_floor}s)" +fi + +summary diff --git a/tools/tests/scenario_crash_restart.sh b/tools/tests/scenario_crash_restart.sh new file mode 100755 index 0000000..869d6c7 --- /dev/null +++ b/tools/tests/scenario_crash_restart.sh @@ -0,0 +1,88 @@ +#!/usr/bin/env bash +# Kills the whole supervisor process group mid-run (SIGKILL, no graceful +# shutdown trap -- simulates a real crash/OOM, not a clean stop) while some +# batches are still in-flight (claimed but not yet committed), then restarts +# pipeline.sh against the SAME EXPORT_ROOT/WATCH_DIR and verifies: the +# restart isn't blocked by the dead process's stale pgid/lock, the stale +# claims get swept and requeued (not stuck forever), nothing that was +# already committed before the kill gets re-exported, and everything ends +# up committed exactly once with a correct final merge. +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" +source ./lib.sh + +TEST_ROOT="/tmp/pipeline_test_crash_restart_$$" +require_test_path "$TEST_ROOT" +trap 'kill -- "-${pid:-}" 2>/dev/null || true; rm -rf "$TEST_ROOT"' EXIT +WATCH="$TEST_ROOT/watch" +EXPORT_ROOT="$TEST_ROOT/export" +mkdir -p "$WATCH" "$EXPORT_ROOT" + +echo "=== scenario_crash_restart ===" + +IDS=(crec01 crec02 crec03 crec04 crec05 crec06) +for id in "${IDS[@]}"; do make_fixture local "$WATCH" "$id" sleep:8; done + +export PIPELINE_MODE=local +export PIPELINE_WATCH_DIR="$WATCH" +export PIPELINE_EXPORT_ROOT="$EXPORT_ROOT" +export PIPELINE_EXPORT_CLI="$FAKE_CLI" +export PIPELINE_EXPORT_CONFIG="$TEST_ROOT/fake_config.json" +touch "$PIPELINE_EXPORT_CONFIG" +export PIPELINE_CHUNK=1 +export PIPELINE_POLL_SECONDS=3 +export PIPELINE_IDLE_MINUTES=1 + +echo "--- first run: will SIGKILL mid-flight ---" +"$PIPELINE" supervisor >"$TEST_ROOT/run1.log" 2>&1 & +pid=$! +# Local mode's 60s candidate-maturity floor means nothing is claimable before +# t=60s; give it a few seconds into the export wave, so some batches are +# genuinely in-flight (claimed, mid sleep:8, not yet committed) when killed. +sleep 66 +old_pgid="$pid" +kill -9 -- "-$pid" 2>/dev/null || true +wait "$pid" 2>/dev/null || true +sleep 1 # let the kernel finish reaping before we look + +committed_before_kill=$(committed_ids "$EXPORT_ROOT" | sort -u) +claimed_before_restart=$(ls "$EXPORT_ROOT/.pipeline/claimed" 2>/dev/null || true) +echo "after kill: committed=[$committed_before_kill] still-claimed=[$claimed_before_restart]" +if [[ -n "$claimed_before_restart" ]]; then + pass "at least one recording was genuinely in-flight (claimed, uncommitted) at kill time" +else + fail "nothing was in-flight at kill time -- test didn't exercise the crash window (timing needs tuning)" +fi +# sanity: nothing should be BOTH committed and still claimed +overlap=$(comm -12 <(sort <<< "$committed_before_kill") <(sort <<< "$claimed_before_restart")) +assert_eq "$overlap" "" "no recording is both committed and still claimed after the kill" + +echo "--- second run: restart against the same EXPORT_ROOT ---" +"$PIPELINE" supervisor >"$TEST_ROOT/run2.log" 2>&1 & +pid=$! +wait "$pid" || true +echo "restart finished" + +if grep -q "still alive" "$TEST_ROOT/run2.log"; then + fail "restart refused to start, claiming the old (dead) process group was still alive" +else + pass "restart was not blocked by the dead process's stale pgid" +fi + +committed_after=$(committed_ids "$EXPORT_ROOT" | sort -u) +for id in "${IDS[@]}"; do + count=$(grep -cx "$id" <<< "$committed_after" || true) + assert_eq "$count" "1" "$id committed exactly once across both runs (no loss, no duplicate)" +done +dupes=$(committed_ids "$EXPORT_ROOT" | sort | uniq -d) +assert_eq "$dupes" "" "no ID committed into more than one batch across both runs" + +merged="${EXPORT_ROOT}_merged" +if [[ -f "$merged/meta/info.json" ]]; then + actual_episodes=$(python3 -c "import json;print(json.load(open('$merged/meta/info.json'))['total_episodes'])") + assert_eq "$actual_episodes" "${#IDS[@]}" "merged dataset episode count after crash+restart" +else + fail "merged dataset exists at $merged" +fi + +summary diff --git a/tools/tests/scenario_local_backlog.sh b/tools/tests/scenario_local_backlog.sh new file mode 100755 index 0000000..1989536 --- /dev/null +++ b/tools/tests/scenario_local_backlog.sh @@ -0,0 +1,94 @@ +#!/usr/bin/env bash +# Local mode, entire backlog present upfront (collection already finished) — +# the "just export, but still use parallelism" case. Also covers: bisection +# isolating a bad recording without quarantining its batch-mates, quarantine +# after 3 permanent failures, and a flaky recording recovering via retry +# instead of being quarantined. Genuine wall-clock concurrency proof is a +# separate test (scenario_concurrency.sh) — batch composition here is +# claim-order-dependent (bisection can group/split slow fixtures +# unpredictably), so a tight timing assertion here would be flaky by +# construction; this scenario instead asserts >1 distinct worker actually +# committed real work, which is deterministic regardless of batch shuffling. +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" +source ./lib.sh + +TEST_ROOT="/tmp/pipeline_test_local_backlog_$$" +require_test_path "$TEST_ROOT" +# pid's process group == pid (pipeline.sh's supervisor re-execs itself under +# setsid before doing anything else), so this reaches the whole acquire/ +# worker/nova-data-cli tree even if this scenario is itself killed/timed out. +trap 'kill -- "-${pid:-}" 2>/dev/null || true; rm -rf "$TEST_ROOT" "$FAKE_CLI_STATE_DIR"' EXIT +export FAKE_CLI_STATE_DIR="$TEST_ROOT/fake_cli_state" +WATCH="$TEST_ROOT/watch" +EXPORT_ROOT="$TEST_ROOT/export" +mkdir -p "$WATCH" "$EXPORT_ROOT" + +echo "=== scenario_local_backlog ===" + +FAST_IDS=(rec01 rec02 rec03 rec04 rec05 rec06 rec07 rec08) +for id in "${FAST_IDS[@]}"; do make_fixture local "$WATCH" "$id" success; done +make_fixture local "$WATCH" bad_always skip # should end up quarantined +make_fixture local "$WATCH" bad_flaky flaky:2 # should recover on 3rd attempt, not quarantined +SLOW_IDS=(slow01 slow02 slow03) +for id in "${SLOW_IDS[@]}"; do make_fixture local "$WATCH" "$id" sleep:5; done + +ALL_IDS=("${FAST_IDS[@]}" bad_always bad_flaky "${SLOW_IDS[@]}") + +export PIPELINE_MODE=local +export PIPELINE_WATCH_DIR="$WATCH" +export PIPELINE_EXPORT_ROOT="$EXPORT_ROOT" +export PIPELINE_EXPORT_CLI="$FAKE_CLI" +export PIPELINE_EXPORT_CONFIG="$TEST_ROOT/fake_config.json" +touch "$PIPELINE_EXPORT_CONFIG" +export PIPELINE_CHUNK=3 +export PIPELINE_POLL_SECONDS=3 +export PIPELINE_IDLE_MINUTES=1 + +"$PIPELINE" supervisor >"$TEST_ROOT/supervisor.log" 2>&1 & +pid=$! +wait "$pid" || true +echo "pipeline finished" + +# --- assertions --- +committed=$(committed_ids "$EXPORT_ROOT" | sort -u) +quarantined=$(quarantined_ids "$EXPORT_ROOT" | sort -u) + +# 1. every fixture ends up committed XOR quarantined, nothing dropped, nothing duplicated +for id in "${ALL_IDS[@]}"; do + count_committed=$(grep -cx "$id" <<< "$committed" || true) + is_quarantined=$(grep -cx "$id" <<< "$quarantined" || true) + if [[ "$id" == "bad_always" ]]; then + assert_eq "$count_committed" "0" "bad_always never committed" + assert_eq "$is_quarantined" "1" "bad_always quarantined" + else + assert_eq "$count_committed" "1" "$id committed exactly once" + assert_eq "$is_quarantined" "0" "$id NOT quarantined" + fi +done + +# 2. no ID appears in more than one batch's .claimed_ids (no duplicate export) +dupes=$(committed_ids "$EXPORT_ROOT" | sort | uniq -d) +assert_eq "$dupes" "" "no ID committed into more than one batch" + +# 3. merge ran exactly once and produced a dataset with the right episode count +expected_episodes=$(( ${#FAST_IDS[@]} + 1 + ${#SLOW_IDS[@]} )) # fast + bad_flaky (recovers) + slow; bad_always contributes 0 +merged="${EXPORT_ROOT}_merged" +if [[ -f "$merged/meta/info.json" ]]; then + actual_episodes=$(python3 -c "import json;print(json.load(open('$merged/meta/info.json'))['total_episodes'])") + assert_eq "$actual_episodes" "$expected_episodes" "merged dataset episode count" +else + fail "merged dataset exists at $merged" +fi + +# 4. more than one worker index actually committed a batch (real parallel use, +# not just WORKERS>1 sitting idle while one worker does everything) +distinct_workers=$(grep -oh 'committed batch_[0-9_]*_w[0-9]*' "$EXPORT_ROOT"/.pipeline/logs/w*.log 2>/dev/null \ + | grep -o '_w[0-9]*$' | sort -u | wc -l) +if [[ $distinct_workers -ge 2 ]]; then + pass "$distinct_workers distinct workers committed batches (real parallel use)" +else + fail "only $distinct_workers distinct worker(s) committed batches" +fi + +summary diff --git a/tools/tests/scenario_local_live.sh b/tools/tests/scenario_local_live.sh new file mode 100755 index 0000000..250e84e --- /dev/null +++ b/tools/tests/scenario_local_live.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# Local mode, recordings trickle in progressively (simulating a collector +# process still running on this machine) instead of all existing upfront — +# verifies exporting actually starts while new recordings are still arriving +# (not after collection "finishes"), and that collection_done doesn't fire +# while the feed is still active. +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" +source ./lib.sh + +TEST_ROOT="/tmp/pipeline_test_local_live_$$" +require_test_path "$TEST_ROOT" +trap 'kill -- "-${pid:-}" 2>/dev/null || true; kill "${collector_pid:-0}" 2>/dev/null || true; rm -rf "$TEST_ROOT"' EXIT +WATCH="$TEST_ROOT/watch" +EXPORT_ROOT="$TEST_ROOT/export" +mkdir -p "$WATCH" "$EXPORT_ROOT" + +echo "=== scenario_local_live ===" + +IDS=(live01 live02 live03 live04 live05 live06) +INTERVAL_S=15 # span (5 gaps * 15s = 75s) must exceed local mode's 60s candidate-maturity + # floor, so the first arrival matures and can be exported *before* + # the collector finishes adding the rest -> genuine overlap. + +# Fake "collector": writes one fixture every INTERVAL_S seconds, then marks itself done. +( + for id in "${IDS[@]}"; do + make_fixture local "$WATCH" "$id" success + sleep "$INTERVAL_S" + done + date +%s > "$TEST_ROOT/collector_finished_at" +) & +collector_pid=$! + +export PIPELINE_MODE=local +export PIPELINE_WATCH_DIR="$WATCH" +export PIPELINE_EXPORT_ROOT="$EXPORT_ROOT" +export PIPELINE_EXPORT_CLI="$FAKE_CLI" +export PIPELINE_EXPORT_CONFIG="$TEST_ROOT/fake_config.json" +touch "$PIPELINE_EXPORT_CONFIG" +export PIPELINE_CHUNK=2 +export PIPELINE_POLL_SECONDS=3 +export PIPELINE_IDLE_MINUTES=1 + +start_epoch=$(date +%s) +"$PIPELINE" supervisor >"$TEST_ROOT/supervisor.log" 2>&1 & +pid=$! +wait "$collector_pid" +wait "$pid" || true +echo "pipeline finished" + +# --- assertions --- +committed=$(committed_ids "$EXPORT_ROOT" | sort -u) +for id in "${IDS[@]}"; do + count=$(grep -cx "$id" <<< "$committed" || true) + assert_eq "$count" "1" "$id committed exactly once" +done + +collector_finished_at=$(cat "$TEST_ROOT/collector_finished_at") +first_commit_at=$(stat -c %Y "$EXPORT_ROOT"/batch_*/dataset 2>/dev/null | sort -n | head -1) +if [[ -n "$first_commit_at" && "$first_commit_at" -lt "$collector_finished_at" ]]; then + pass "first export committed ($((first_commit_at - start_epoch))s in) before collector finished ($((collector_finished_at - start_epoch))s in) -> pull+export overlapped" +else + fail "first export ($first_commit_at) did not precede collector finishing ($collector_finished_at) -- no overlap observed" +fi + +collection_done_at=$(stat -c %Y "$EXPORT_ROOT/.pipeline/collection_done" 2>/dev/null || echo 0) +if [[ "$collection_done_at" -gt "$collector_finished_at" ]]; then + pass "collection_done ($collection_done_at) fired after collector finished ($collector_finished_at), not before" +else + fail "collection_done ($collection_done_at) fired at/before collector finished ($collector_finished_at) -- premature idle detection" +fi + +merged="${EXPORT_ROOT}_merged" +if [[ -f "$merged/meta/info.json" ]]; then + actual_episodes=$(python3 -c "import json;print(json.load(open('$merged/meta/info.json'))['total_episodes'])") + assert_eq "$actual_episodes" "${#IDS[@]}" "merged dataset episode count" +else + fail "merged dataset exists at $merged" +fi + +summary diff --git a/tools/tests/scenario_real_smoke.sh b/tools/tests/scenario_real_smoke.sh new file mode 100755 index 0000000..aa893d9 --- /dev/null +++ b/tools/tests/scenario_real_smoke.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash +# Tier 2: ONE real end-to-end run through the actual nova-data-cli (no stub), +# using two small pre-existing sample recordings, copied (never symlinked or +# moved) from /home/sebi/ws/Data/5-pick-cube-sim-raw into an isolated test +# dir -- catches anything the fake_nova_data_cli.py stub can't: real +# export_summary.json shape, real aggregate_datasets behavior on a real +# dataset, real timing. Local mode only (remote transport is already covered +# for real in scenario_remote_*.sh with the stub; this test's job is the real +# CLI, not re-proving rsync). +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" +source ./lib.sh + +SRC_ROOT="/home/sebi/ws/Data/5-pick-cube-sim-raw" +REAL_CONFIG="/home/sebi/ws/pick_and_place_imitation_learning/data_collection/configs/lerobot_export.json" + +TEST_ROOT="/tmp/pipeline_test_real_smoke_$$" +require_test_path "$TEST_ROOT" +trap 'kill -- "-${pid:-}" 2>/dev/null || true; rm -rf "$TEST_ROOT"' EXIT +WATCH="$TEST_ROOT/watch" +EXPORT_ROOT="$TEST_ROOT/export" +mkdir -p "$WATCH" "$EXPORT_ROOT" + +echo "=== scenario_real_smoke (real nova-data-cli, no stub) ===" + +if [[ ! -d "$SRC_ROOT" || ! -f "$REAL_CONFIG" ]]; then + echo "SKIP: sample recordings ($SRC_ROOT) or real config ($REAL_CONFIG) not found on this machine" + exit 0 +fi + +# Two smallest available recordings, to keep this test's runtime reasonable. +mapfile -t SAMPLE_IDS < <( + for d in "$SRC_ROOT"/*/; do + [[ -f "$d/recording.rrd" ]] || continue + echo "$(stat -c %s "$d/recording.rrd") $(basename "$d")" + done | sort -n | head -2 | awk '{print $2}' +) +[[ ${#SAMPLE_IDS[@]} -eq 2 ]] || { echo "SKIP: fewer than 2 sample recordings with recording.rrd found"; exit 0; } + +for id in "${SAMPLE_IDS[@]}"; do + cp -r "$SRC_ROOT/$id" "$WATCH/$id" # copy, never touch the source +done + +export PIPELINE_MODE=local +export PIPELINE_WATCH_DIR="$WATCH" +export PIPELINE_EXPORT_ROOT="$EXPORT_ROOT" +export PIPELINE_EXPORT_CONFIG="$REAL_CONFIG" +# PIPELINE_EXPORT_CLI left unset -> real `uv run nova-data-cli` +export PIPELINE_CHUNK=2 +export PIPELINE_POLL_SECONDS=5 +export PIPELINE_IDLE_MINUTES=1 + +echo "exporting ${SAMPLE_IDS[*]} through the real CLI (this actually decodes/encodes, expect ~1-3 min)..." +"$PIPELINE" supervisor >"$TEST_ROOT/supervisor.log" 2>&1 & +pid=$! +wait "$pid" || true +echo "pipeline finished" + +committed=$(committed_ids "$EXPORT_ROOT" | sort -u) +for id in "${SAMPLE_IDS[@]}"; do + count=$(grep -cx "$id" <<< "$committed" || true) + assert_eq "$count" "1" "$id committed exactly once (real CLI)" +done + +batch_dir=$(find "$EXPORT_ROOT" -maxdepth 1 -name 'batch_*' | head -1) +if [[ -n "$batch_dir" && -f "$batch_dir/dataset/export_summary.json" ]]; then + pass "real export_summary.json written: $(cat "$batch_dir/dataset/export_summary.json")" +else + fail "no real export_summary.json found in committed batch" +fi + +merged="${EXPORT_ROOT}_merged" +if [[ -f "$merged/meta/info.json" ]]; then + actual_episodes=$(python3 -c "import json;print(json.load(open('$merged/meta/info.json'))['total_episodes'])") + if [[ "$actual_episodes" -ge 1 ]]; then + pass "real merged dataset produced ($actual_episodes episode(s))" + else + fail "real merged dataset has 0 episodes" + fi +else + fail "merged dataset exists at $merged" +fi + +summary diff --git a/tools/tests/scenario_remote_backlog.sh b/tools/tests/scenario_remote_backlog.sh new file mode 100755 index 0000000..d1890ee --- /dev/null +++ b/tools/tests/scenario_remote_backlog.sh @@ -0,0 +1,68 @@ +#!/usr/bin/env bash +# Remote mode (real SSH/rsync against the real workstation host), entire +# backlog present upfront on the remote side. Uses a dedicated, brand-new +# subdir under /mnt/data/sebastian/pipeline_test_fixtures/ on the real remote +# host -- NEVER the real raw_datasets/ collection dir. +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" +source ./lib.sh + +RUN_ID="run_backlog_$$" +REMOTE_DIR="$REMOTE_TEST_ROOT/$RUN_ID" +require_test_path "$REMOTE_DIR" # contains "pipeline_test" -- extra guard before any remote rm -rf + +TEST_ROOT="/tmp/pipeline_test_remote_backlog_$$" +require_test_path "$TEST_ROOT" +cleanup() { + kill -- "-${pid:-}" 2>/dev/null || true + ssh -o BatchMode=yes "$REMOTE_TEST_HOST" "rm -rf '$REMOTE_DIR'" 2>/dev/null || true + rm -rf "$TEST_ROOT" +} +trap cleanup EXIT +mkdir -p "$TEST_ROOT/watch" "$TEST_ROOT/export" + +echo "=== scenario_remote_backlog ===" + +IDS=(rrec01 rrec02 rrec03 rrec04 rrec05) +for id in "${IDS[@]}"; do make_fixture remote "$REMOTE_DIR" "$id" success; done +make_fixture remote "$REMOTE_DIR" rbad skip # should end up quarantined + +export PIPELINE_MODE=remote +export PIPELINE_REMOTE_HOST="$REMOTE_TEST_HOST" +export PIPELINE_REMOTE_DIRS="$REMOTE_DIR" +export PIPELINE_WATCH_DIR="$TEST_ROOT/watch/placeholder" # only its dirname is used; basename comes from REMOTE_DIRS +export PIPELINE_EXPORT_ROOT="$TEST_ROOT/export" +export PIPELINE_EXPORT_CLI="$FAKE_CLI" +export PIPELINE_EXPORT_CONFIG="$TEST_ROOT/fake_config.json" +touch "$PIPELINE_EXPORT_CONFIG" +export PIPELINE_CHUNK=3 +export PIPELINE_POLL_SECONDS=5 +export PIPELINE_IDLE_MINUTES=1 + +"$PIPELINE" supervisor >"$TEST_ROOT/supervisor.log" 2>&1 & +pid=$! +wait "$pid" || true +echo "pipeline finished" + +WATCH_ACTUAL="$TEST_ROOT/watch/$RUN_ID" +[[ -d "$WATCH_ACTUAL" ]] && pass "rsync pulled recordings into local watch dir" \ + || fail "rsync did not create expected local watch dir $WATCH_ACTUAL" + +committed=$(committed_ids "$TEST_ROOT/export" | sort -u) +quarantined=$(quarantined_ids "$TEST_ROOT/export" | sort -u) +for id in "${IDS[@]}"; do + count=$(grep -cx "$id" <<< "$committed" || true) + assert_eq "$count" "1" "$id committed exactly once" +done +is_quarantined=$(grep -cx "rbad" <<< "$quarantined" || true) +assert_eq "$is_quarantined" "1" "rbad quarantined" + +merged="${TEST_ROOT}/export_merged" +if [[ -f "$merged/meta/info.json" ]]; then + actual_episodes=$(python3 -c "import json;print(json.load(open('$merged/meta/info.json'))['total_episodes'])") + assert_eq "$actual_episodes" "${#IDS[@]}" "merged dataset episode count" +else + fail "merged dataset exists at $merged" +fi + +summary diff --git a/tools/tests/scenario_remote_live.sh b/tools/tests/scenario_remote_live.sh new file mode 100755 index 0000000..0a08138 --- /dev/null +++ b/tools/tests/scenario_remote_live.sh @@ -0,0 +1,87 @@ +#!/usr/bin/env bash +# Remote mode (real SSH/rsync), recordings arrive progressively on the remote +# side while the pipeline is already running -- verifies pull+export overlap +# and that idle-detection doesn't fire while the remote feed is still active. +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" +source ./lib.sh + +RUN_ID="run_live_$$" +REMOTE_DIR="$REMOTE_TEST_ROOT/$RUN_ID" +require_test_path "$REMOTE_DIR" + +TEST_ROOT="/tmp/pipeline_test_remote_live_$$" +require_test_path "$TEST_ROOT" +cleanup() { + kill -- "-${pid:-}" 2>/dev/null || true + kill "${collector_pid:-0}" 2>/dev/null || true + ssh -o BatchMode=yes "$REMOTE_TEST_HOST" "rm -rf '$REMOTE_DIR'" 2>/dev/null || true + rm -rf "$TEST_ROOT" +} +trap cleanup EXIT +mkdir -p "$TEST_ROOT/watch" "$TEST_ROOT/export" + +echo "=== scenario_remote_live ===" + +IDS=(rlive01 rlive02 rlive03 rlive04) +INTERVAL_S=12 # remote mode has no candidate-maturity floor (unlike local), so + # this just needs to comfortably outlast a couple of rsync polls + +( + for id in "${IDS[@]}"; do + make_fixture remote "$REMOTE_DIR" "$id" success + sleep "$INTERVAL_S" + done + date +%s > "$TEST_ROOT/collector_finished_at" +) & +collector_pid=$! + +export PIPELINE_MODE=remote +export PIPELINE_REMOTE_HOST="$REMOTE_TEST_HOST" +export PIPELINE_REMOTE_DIRS="$REMOTE_DIR" +export PIPELINE_WATCH_DIR="$TEST_ROOT/watch/placeholder" +export PIPELINE_EXPORT_ROOT="$TEST_ROOT/export" +export PIPELINE_EXPORT_CLI="$FAKE_CLI" +export PIPELINE_EXPORT_CONFIG="$TEST_ROOT/fake_config.json" +touch "$PIPELINE_EXPORT_CONFIG" +export PIPELINE_CHUNK=1 +export PIPELINE_POLL_SECONDS=4 +export PIPELINE_IDLE_MINUTES=1 + +start_epoch=$(date +%s) +"$PIPELINE" supervisor >"$TEST_ROOT/supervisor.log" 2>&1 & +pid=$! +wait "$collector_pid" +wait "$pid" || true +echo "pipeline finished" + +committed=$(committed_ids "$TEST_ROOT/export" | sort -u) +for id in "${IDS[@]}"; do + count=$(grep -cx "$id" <<< "$committed" || true) + assert_eq "$count" "1" "$id committed exactly once" +done + +collector_finished_at=$(cat "$TEST_ROOT/collector_finished_at") +first_commit_at=$(stat -c %Y "$TEST_ROOT"/export/batch_*/dataset 2>/dev/null | sort -n | head -1) +if [[ -n "$first_commit_at" && "$first_commit_at" -lt "$collector_finished_at" ]]; then + pass "first export committed ($((first_commit_at - start_epoch))s in) before remote collector finished ($((collector_finished_at - start_epoch))s in) -> pull+export overlapped" +else + fail "first export ($first_commit_at) did not precede collector finishing ($collector_finished_at) -- no overlap observed" +fi + +collection_done_at=$(stat -c %Y "$TEST_ROOT/export/.pipeline/collection_done" 2>/dev/null || echo 0) +if [[ "$collection_done_at" -gt "$collector_finished_at" ]]; then + pass "collection_done fired after remote collector finished, not before" +else + fail "collection_done ($collection_done_at) fired at/before collector finished ($collector_finished_at) -- premature idle detection" +fi + +merged="${TEST_ROOT}/export_merged" +if [[ -f "$merged/meta/info.json" ]]; then + actual_episodes=$(python3 -c "import json;print(json.load(open('$merged/meta/info.json'))['total_episodes'])") + assert_eq "$actual_episodes" "${#IDS[@]}" "merged dataset episode count" +else + fail "merged dataset exists at $merged" +fi + +summary diff --git a/tools/tests/scenario_sizing.sh b/tools/tests/scenario_sizing.sh new file mode 100755 index 0000000..e661537 --- /dev/null +++ b/tools/tests/scenario_sizing.sh @@ -0,0 +1,50 @@ +#!/usr/bin/env bash +# Unit-level check of compute_workers() via pipeline.sh's `sizing` debug role +# (case dispatch hook added for exactly this) -- no fixtures, no export, just +# verifies WORKERS scales with MEM_TARGET_FRACTION/WORKER_MEM_ESTIMATE_MB and +# clamps to >=1 and to a fraction of nproc, using this real machine's actual +# /proc/meminfo and nproc (no mocking needed -- the formula is pure arithmetic +# over real values, so this is a faithful check without a fixture harness). +set -euo pipefail +cd "$(dirname "${BASH_SOURCE[0]}")" +source ./lib.sh + +TEST_ROOT="/tmp/pipeline_test_sizing_$$" +require_test_path "$TEST_ROOT" +trap 'rm -rf "$TEST_ROOT"' EXIT +mkdir -p "$TEST_ROOT/watch" "$TEST_ROOT/export" + +echo "=== scenario_sizing ===" + +mem_total_kb=$(awk '/MemTotal/{print $2}' /proc/meminfo) +mem_total_mb=$((mem_total_kb / 1024)) +cores=$(nproc) + +sizing() { + PIPELINE_MODE=local PIPELINE_WATCH_DIR="$TEST_ROOT/watch" PIPELINE_EXPORT_ROOT="$TEST_ROOT/export" \ + PIPELINE_MEM_TARGET_FRACTION="$1" PIPELINE_WORKER_MEM_ESTIMATE_MB="$2" \ + "$PIPELINE" sizing --mode local +} + +# 1. a tiny fraction / huge per-worker estimate should clamp to the floor of 1 +w=$(sizing 0.01 999999999) +assert_eq "$w" "1" "clamps to minimum of 1 worker when budget is far below one worker's estimate" + +# 2. an unrealistically low per-worker estimate should clamp to the nproc-derived ceiling +w=$(sizing 0.99 1) +core_ceiling=$(( cores * 8 / 10 )) +assert_eq "$w" "$core_ceiling" "clamps to the nproc-derived ceiling (80% of $cores cores) when memory allows far more" + +# 3. doubling MEM_TARGET_FRACTION (while staying under the core ceiling) should +# roughly double the worker count -- proves it actually scales with the +# fraction instead of being some other hardcoded number +small_est=$((mem_total_mb / 20)) # deliberately large estimate so both fractions stay core-ceiling-safe +w_low=$(sizing 0.10 "$small_est") +w_high=$(sizing 0.20 "$small_est") +if [[ $w_high -ge $((w_low * 2 - 1)) && $w_high -le $((w_low * 2 + 1)) ]]; then + pass "worker count scales with MEM_TARGET_FRACTION ($w_low -> $w_high for 0.10 -> 0.20)" +else + fail "worker count did not scale as expected ($w_low -> $w_high for 0.10 -> 0.20)" +fi + +summary diff --git a/tools/validate_batch.py b/tools/validate_batch.py new file mode 100755 index 0000000..3605922 --- /dev/null +++ b/tools/validate_batch.py @@ -0,0 +1,57 @@ +#!/usr/bin/env python +"""Validate a pipeline.sh batch export before it's committed. + +Checks export_summary.json against the batch's .claimed_ids. At batch size 1, +a skip/fail is attributable to that one recording; larger batches are only +checked in aggregate — see AGENT.md for why, and how the caller bisects. +""" + +import argparse +import json +import sys +from pathlib import Path + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output-dir", required=True, type=Path, help="Batch's tmp output dir") + parser.add_argument("--claimed-ids", required=True, type=Path, help="Path to the batch's .claimed_ids file") + args = parser.parse_args() + + claimed = [line.strip() for line in args.claimed_ids.read_text().splitlines() if line.strip()] + if not claimed: + print("No claimed IDs — nothing to validate", file=sys.stderr) + return 1 + + summary_path = args.output_dir / "export_summary.json" + if not summary_path.is_file(): + print(f"No export_summary.json in {args.output_dir}", file=sys.stderr) + return 1 + + summary = json.loads(summary_path.read_text()) + successful = summary.get("successful_episodes", 0) + skipped = summary.get("skipped_episodes", 0) + failed = summary.get("failed_episodes", 0) + + if len(claimed) == 1: + if successful > 0: + print(f"OK: {claimed[0]} exported ({successful} episode(s))") + return 0 + print(f"FAIL: {claimed[0]} produced no episodes (skipped={skipped}, failed={failed})", file=sys.stderr) + return 1 + + if skipped == 0 and failed == 0: + print(f"OK: all {len(claimed)} claimed recordings exported cleanly") + return 0 + + print( + f"AMBIGUOUS: batch of {len(claimed)} had {skipped} skipped / {failed} failed episodes, " + "cannot attribute to a specific recording — bisect required. " + f"Claimed: {', '.join(claimed)}", + file=sys.stderr, + ) + return 1 + + +if __name__ == "__main__": + sys.exit(main()) From 9ea1dab80da5f5f5bf2abcf9b7add000125d8004 Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Thu, 6 Aug 2026 17:59:38 +0200 Subject: [PATCH 04/12] Require repeated confirmation before treating a scan as "nothing left" role_worker exited for good on a single empty candidate scan once COLLECTION_DONE existed, with no tolerance for a transient glitch in that one scan. Observed in production: worker w3 exited claiming no candidates while ~800 unclaimed recordings were still sitting right there, apparently under system load (4 CPU-heavy nova-data-cli decodes competing for cores) - a single false-empty scan permanently dropped a quarter of the throughput with no error signal. role_supervisor's drain loop had the same shape and a worse failure mode (would merge before everything was actually exported). Both now require repeated confirmation (3 consecutive empty scans for a worker, 2 clean passes for the supervisor) before trusting "nothing left", matching the pattern role_acquire's idle-detection already used. Also adds a final report (recordings found/exported/quarantined, episodes in the merged dataset) and a nonzero exit code when anything's quarantined, so a systemic failure doesn't silently look like a clean run. Co-Authored-By: Claude Sonnet 5 --- docs/investigations/worker-early-exit.md | 66 ++++++++++++++++++++++++ tools/AGENT.md | 17 ++++++ tools/pipeline.sh | 64 ++++++++++++++++++++--- 3 files changed, 141 insertions(+), 6 deletions(-) create mode 100644 docs/investigations/worker-early-exit.md diff --git a/docs/investigations/worker-early-exit.md b/docs/investigations/worker-early-exit.md new file mode 100644 index 0000000..36ff4f2 --- /dev/null +++ b/docs/investigations/worker-early-exit.md @@ -0,0 +1,66 @@ +# Symptom: worker exits early, permanently down one worker + +## Observed + +During a `tools/pipeline.sh` run exporting the `choreo2` dataset +(`pick_and_place_sim_20260805_191245`, 1001 recordings, `--mode local`, +started 2026-08-06 ~17:07), worker `w3` logged: + +``` +[w3] [17:17:58] [worker] committed batch_20260806_171301_w3_1147218_10206 (8 recording(s)) +[w3] [17:17:59] [worker] no candidates and collection done, exiting +``` + +...and never restarted. From then on, only `w0`, `w1`, `w2` were running +(confirmed via `ps aux | grep nova-data-cli`), each still committing batches +every ~5 minutes. + +## Why it doesn't add up + +At 17:17:59: +- ~160 episodes committed (21 batches × 8, from the `w0-w3` logs) +- At most ~32 in flight (4 workers × `CHUNK=8`) +- That leaves **~800 unclaimed candidates** still sitting in + `/home/sebi/ws/Data/raw_data/choreo2/pick_and_place_sim_20260805_191245` + +`collection_done` was legitimately set at 17:07:56 (the earlier remote-mode +acquire pass had already finished its final rsync at 17:01:36, so all 1001 +recordings were genuinely present on disk by the time the local-mode watcher +started). So the "collection is done" half of the exit condition is correct. + +The "no candidates" half is not plausible given ~800 unclaimed recordings +should have been sitting right there for `list_candidates()` to find. + +## Effect + +- One of four workers permanently disappears mid-run, with no log line + indicating an error — just the same message a worker prints on genuine, + correct completion. +- The remaining three workers keep making progress, so nothing looks wrong + at a glance (batches keep committing), but overall throughput drops to + 75% for the rest of the run. +- Nothing in the pipeline's own signals (logs, exit code, `.pipeline/` + state) distinguishes this from the expected end-of-run shutdown. + +## Reproduction notes + +Not reproduced in isolation — several attempts to trigger the same premature +`set -e`-style abort or scan failure in `list_candidates`/`is_candidate` +(including under a synthetic 50-recording watch dir with the exact same +functions) came back clean, so the underlying trigger is genuinely +load-dependent, not a deterministic logic bug in the scan itself. + +## Fix + +`role_worker`'s exit condition trusted a *single* empty scan as proof there +was no more work, the moment `COLLECTION_DONE` was also true — with no +tolerance for a transient glitch in that one scan (e.g. a `find` fork failing +under the exact system load described above). `role_supervisor`'s drain loop +had the same shape and a worse consequence (a false-empty scan there would +merge before everything was actually exported, not just drop a worker). + +Both now require repeated confirmation before trusting "nothing left" — 3 +consecutive empty scans for a worker to exit, 2 consecutive clean passes for +the supervisor to proceed to merge — the same pattern `role_acquire`'s +idle-detection already used and this code didn't. See `tools/pipeline.sh` +`role_worker`/`role_supervisor` and `tools/AGENT.md`. diff --git a/tools/AGENT.md b/tools/AGENT.md index 0dfe011..9b72e88 100644 --- a/tools/AGENT.md +++ b/tools/AGENT.md @@ -152,6 +152,23 @@ keeps this from monopolizing CPU/disk even without an active watchdog for those resources — it tells the kernel to prefer any other process, so the pipeline only consumes spare capacity. +## "Nothing left" needs confirmation, not a single scan + +A worker exits when it finds no candidates *and* `COLLECTION_DONE` exists; the +supervisor's drain loop proceeds to merge on the same condition. Both used to +trust a single scan of `list_candidates` for this — but a scan can come back +empty transiently (observed in production under heavy system load, likely a +`find` subprocess failing to fork while several CPU-heavy `nova-data-cli` +decodes were competing for cores; see +`docs/investigations/worker-early-exit.md`), and there was no tolerance for +that before treating it as final. A false-empty scan for a worker silently +drops it for the rest of the run (throughput loss, no error); the same false +reading in the supervisor's drain loop would be worse — merging before +everything's actually exported. Both now require repeated confirmation (3 +consecutive empty scans for a worker, 2 consecutive clean passes for the +supervisor) before trusting it, the same principle `role_acquire`'s +idle-detection already applied. + ## Merge `merge_batches.py` merges via `lerobot.datasets.aggregate.aggregate_datasets`, diff --git a/tools/pipeline.sh b/tools/pipeline.sh index afdfdb5..ac2ab7c 100755 --- a/tools/pipeline.sh +++ b/tools/pipeline.sh @@ -272,6 +272,7 @@ attempt_batch() { role_worker() { WORKER_IDX="$1" + local empty_scans=0 while true; do rebuild_done wait_for_memory @@ -299,14 +300,23 @@ role_worker() { fi if [[ ${#claimed[@]} -eq 0 ]]; then - if [[ -f "$COLLECTION_DONE" ]]; then - log "no candidates and collection done, exiting" + # A single empty scan can be a transient glitch (e.g. under heavy + # system load), not proof there's no more work — require 3 consecutive + # empty scans before treating "collection done" as "exit for good", + # the same way acquire's idle-detection needs sustained evidence rather + # than a single reading. A false-empty scan just wastes 30s here; a + # false-permanent worker exit silently cuts throughput for the rest of + # the run. + empty_scans=$((empty_scans + 1)) + if [[ -f "$COLLECTION_DONE" && $empty_scans -ge 3 ]]; then + log "no candidates across 3 consecutive scans and collection done, exiting" return 0 fi sleep 30 continue fi + empty_scans=0 attempt_batch "${claimed[@]}" done } @@ -362,25 +372,67 @@ role_supervisor() { wait "$acquire_pid" || true wait "${worker_pids[@]}" || true - # Drain: relaunch until a pass finds nothing left (see AGENT.md). - while true; do + # Drain: relaunch until 2 consecutive passes find nothing left (see + # AGENT.md) — a single clean scan could be a transient glitch, and trusting + # it alone here risks merging before everything's actually exported, not + # just losing a worker the way role_worker's equivalent check would. + local drained_scans=0 + while [[ $drained_scans -lt 2 ]]; do rebuild_done if [[ -z "$(list_candidates | head -1)" ]] && [[ -z "$(ls -A "$CLAIMED" 2>/dev/null)" ]]; then - break + drained_scans=$((drained_scans + 1)) + sleep 5 + continue fi + drained_scans=0 log "drain pass found leftover work, relaunching workers" spawn_workers wait "${worker_pids[@]}" || true done local -a batch_dirs=("$EXPORT_ROOT"/batch_*) - [[ ${#batch_dirs[@]} -gt 0 ]] || { log "no batches produced, nothing to merge"; return 0; } + if [[ ${#batch_dirs[@]} -eq 0 ]]; then + log "no batches produced, nothing to merge" + print_report + [[ -n "$(ls -A "$QUARANTINE" 2>/dev/null)" ]] && exit 1 + return 0 + fi log "all workers drained, merging" (cd "$NOVA_CLI_DIR" && uv run python tools/merge_batches.py \ --batches-root "$EXPORT_ROOT" \ --output "${EXPORT_ROOT}_merged") log "done: ${EXPORT_ROOT}_merged" + print_report + [[ -n "$(ls -A "$QUARANTINE" 2>/dev/null)" ]] && exit 1 + return 0 +} + +# recordings found vs. exported vs. quarantined, and episodes in the merged +# dataset (not the same number — one recording can yield several episodes). +print_report() { + local total=0 + local d + for d in "$WATCH_DIR"/*/; do + [[ -f "${d}recording.rrd" ]] && total=$((total + 1)) + done + rebuild_done + local exported=${#DONE_IDS[@]} + local -a quarantined_ids=("$QUARANTINE"/*) + local quarantined_n=${#quarantined_ids[@]} + + log "=== report ===" + log "recordings found: $total" + log "exported: $exported" + log "quarantined (3x fail): $quarantined_n" + if [[ $quarantined_n -gt 0 ]]; then + log "quarantined IDs: $(basename -a "${quarantined_ids[@]}" | tr '\n' ' ')" + fi + if [[ -f "${EXPORT_ROOT}_merged/meta/info.json" ]]; then + local episodes + episodes="$(python3 -c "import json;print(json.load(open('${EXPORT_ROOT}_merged/meta/info.json'))['total_episodes'])" 2>/dev/null || echo "?")" + log "episodes in merged dataset: $episodes" + fi } case "$ROLE" in From 83f8d875ac2c06da6bc242bce59bba8550ccdb28 Mon Sep 17 00:00:00 2001 From: spereira02 Date: Mon, 10 Aug 2026 19:55:48 +0200 Subject: [PATCH 05/12] feat: added metadata to export --- src/nova_export/export/config.py | 12 ++ src/nova_export/export/episode_sampler.py | 1 + src/nova_export/export/exporter.py | 50 ++++++++ src/nova_export/export/heads/lerobot.py | 36 ++++++ tools/pipeline.sh | 138 +++++++++++++++------- 5 files changed, 194 insertions(+), 43 deletions(-) diff --git a/src/nova_export/export/config.py b/src/nova_export/export/config.py index dbd8428..aa0c78f 100644 --- a/src/nova_export/export/config.py +++ b/src/nova_export/export/config.py @@ -163,6 +163,18 @@ class ExportConfig(BaseModel): description="Dataset identifier (repo_id for LeRobot, dataset name for Groot)", ) + episode_metadata: list[str] = Field( + default_factory=list, + description=( + "meta.json field names (e.g. 'cube_x_mm', 'cube_y_mm', 'cube_yaw_rad') " + "to add as extra columns on each episode's row in meta/episodes/*.parquet " + "(one value per episode, not repeated per frame). Requires local " + "rrd_paths export — meta.json must sit next to each recording's " + "recording.rrd. Ignored (with a warning) when exporting from a remote " + "catalog_url, since there's no local meta.json to read." + ), + ) + @field_validator("index_column") @classmethod def _validate_index_column(cls, v: str) -> str: diff --git a/src/nova_export/export/episode_sampler.py b/src/nova_export/export/episode_sampler.py index 354a6fc..71eacf7 100644 --- a/src/nova_export/export/episode_sampler.py +++ b/src/nova_export/export/episode_sampler.py @@ -63,6 +63,7 @@ class Episode: segment_id: str episode_index: int samples: list[Sample] + extra_metadata: dict[str, float] | None = None @property def num_frames(self) -> int: diff --git a/src/nova_export/export/exporter.py b/src/nova_export/export/exporter.py index bfddba8..ec008ea 100644 --- a/src/nova_export/export/exporter.py +++ b/src/nova_export/export/exporter.py @@ -20,6 +20,7 @@ from __future__ import annotations import contextlib +import json from collections.abc import Callable, Generator from pathlib import Path @@ -199,6 +200,37 @@ def _validate_sources(dataset, config: ExportConfig, segment_id: str) -> None: ) +def _load_episode_metadata( + rrd_paths: list[Path] | None, fields: list[str] +) -> dict[str, dict[str, float]]: + """Read episode_metadata fields from each recording's sibling meta.json. + + Keyed by segment_id, which for local rrd_paths exports is exactly the + recording's directory name (//recording.rrd) — + the same recording_id the collector assigns and rerun uses as the + segment ID, so no separate ID plumbing is needed. + """ + if not fields: + return {} + if not rrd_paths: + logger.warning( + "episode_metadata {} configured but exporting from catalog_url " + "(no local meta.json available) — skipping", + fields, + ) + return {} + + result: dict[str, dict[str, float]] = {} + for rrd_path in rrd_paths: + meta_path = rrd_path.parent / "meta.json" + if not meta_path.is_file(): + continue + meta = json.loads(meta_path.read_text()) + segment_id = rrd_path.parent.name + result[segment_id] = {f: meta[f] for f in fields if f in meta} + return result + + def _raise_fd_limit_for(num_files: int) -> None: """Best-effort: raise this process's open-file limit to fit num_files. @@ -360,6 +392,10 @@ def export_recordings( # Create the export head (Layer 2: format-specific writer) head = _create_export_head(config, output_dir) + episode_metadata_by_segment = _load_episode_metadata( + rrd_paths, config.episode_metadata + ) + # For the max_episode_duration_s safety check: fetch the dataset's # per-segment raw time ranges once (a cheap manifest-metadata read), # rather than once per episode. @@ -452,6 +488,20 @@ def export_recordings( ) continue + if config.episode_metadata: + found = episode_metadata_by_segment.get(segment_id, {}) + missing = [f for f in config.episode_metadata if f not in found] + if missing: + logger.warning( + "Episode {} ({}): meta.json missing {} — filled with 0.0", + episode_id, + segment_id[:8], + missing, + ) + episode.extra_metadata = { + f: found.get(f, 0.0) for f in config.episode_metadata + } + # Check if episode has samples if not episode.samples: reason = "No samples in episode" diff --git a/src/nova_export/export/heads/lerobot.py b/src/nova_export/export/heads/lerobot.py index 15faf4d..c242946 100644 --- a/src/nova_export/export/heads/lerobot.py +++ b/src/nova_export/export/heads/lerobot.py @@ -40,6 +40,7 @@ def __init__(self, config: ExportConfig, output_dir: Path): super().__init__(config, output_dir) self._dataset: LeRobotDataset | None = None self._features: dict[str, Any] | None = None + self._episode_metadata_by_index: dict[int, dict[str, float]] = {} @property def format_name(self) -> str: @@ -147,6 +148,15 @@ def write_episode(self, episode: Episode) -> bool: episode.duration_s, ) + if self.config.episode_metadata and episode.extra_metadata: + # LeRobot assigns its own sequential episode_index (meta.total_episodes) + # when save_episode() runs, which is NOT episode.episode_index (that's + # the exporter's raw segment-loop counter, and diverges as soon as any + # earlier segment is skipped). Key by the index LeRobot is about to use. + self._episode_metadata_by_index[self._dataset.meta.total_episodes] = ( + episode.extra_metadata + ) + try: for sample in tqdm(episode.samples, desc="Frames", leave=False): frame = self._sample_to_frame(sample) @@ -172,6 +182,9 @@ def finalize(self) -> ExportResult: logger.info("Finalizing LeRobot dataset...") self._dataset.finalize() + if self.config.episode_metadata: + self._write_episode_metadata_columns() + logger.success( "Dataset finalized: {} episodes, {} frames → {}", self._dataset.num_episodes, @@ -191,6 +204,29 @@ def finalize(self) -> ExportResult: }, ) + def _write_episode_metadata_columns(self) -> None: + """Add config.episode_metadata fields as columns to meta/episodes/*.parquet. + + One row per episode already exists there (episode_index, length, tasks, + stats, ...) — this just adds our extra columns to those existing rows, + rather than duplicating the values onto every per-frame row in data/. + """ + import pandas as pd + + episodes_files = sorted(self.output_dir.glob("meta/episodes/**/*.parquet")) + for path in episodes_files: + df = pd.read_parquet(path) + for field in self.config.episode_metadata: + df[field] = df["episode_index"].map( + lambda ep: self._episode_metadata_by_index.get(ep, {}).get(field) + ) + df.to_parquet(path) + logger.info( + "Added episode_metadata columns {} to {} episode metadata file(s)", + self.config.episode_metadata, + len(episodes_files), + ) + def _sample_to_frame(self, sample: Sample) -> dict[str, Any]: """Convert a Sample to a LeRobot frame dict. diff --git a/tools/pipeline.sh b/tools/pipeline.sh index ac2ab7c..1b98dfd 100755 --- a/tools/pipeline.sh +++ b/tools/pipeline.sh @@ -7,20 +7,20 @@ shopt -s nullglob # ---- config ----------------------------------------------------------- # Every value is overridable via a PIPELINE_* env var; see README.md. -MODE="${PIPELINE_MODE:-remote}" +MODE="${PIPELINE_MODE:-local}" REMOTE_HOST="${PIPELINE_REMOTE_HOST:-intern@172.31.11.129}" if [[ -n "${PIPELINE_REMOTE_DIRS:-}" ]]; then IFS=':' read -r -a REMOTE_DIRS <<< "$PIPELINE_REMOTE_DIRS" else REMOTE_DIRS=( - "/mnt/data/sebastian/raw_datasets/pick_and_place_sim_20260805_191245" + "/mnt/data/sebastian/raw_datasets/dryrun_pose_algo_check" ) fi -WATCH_DIR="${PIPELINE_WATCH_DIR:-/home/sebi/ws/Data/raw_data/choreo2/pick_and_place_sim_20260805_191245}" +WATCH_DIR="${PIPELINE_WATCH_DIR:-/mnt/data/sebastian/raw_datasets/dryrun_pose_algo_check}" -NOVA_CLI_DIR="${PIPELINE_NOVA_CLI_DIR:-/home/sebi/ws/nova-data-cli}" -EXPORT_CONFIG="${PIPELINE_EXPORT_CONFIG:-/home/sebi/ws/pick_and_place_imitation_learning/data_collection/configs/lerobot_export.json}" -EXPORT_ROOT="${PIPELINE_EXPORT_ROOT:-/home/sebi/ws/Data/choreo2_export}" +NOVA_CLI_DIR="${PIPELINE_NOVA_CLI_DIR:-/home/intern/ws/nova-data-cli}" +EXPORT_CONFIG="${PIPELINE_EXPORT_CONFIG:-/home/intern/ws/pick_and_place_imitation_learning/data_collection/configs/lerobot_export.json}" +EXPORT_ROOT="${PIPELINE_EXPORT_ROOT:-/mnt/data/sebastian/lerobot_datasets/dryrun_pose_algo_check}" read -r -a EXPORT_CLI_CMD <<< "${PIPELINE_EXPORT_CLI:-uv run nova-data-cli}" # swap in a stub for tests CHUNK="${PIPELINE_CHUNK:-8}" @@ -29,8 +29,13 @@ IDLE_MINUTES="${PIPELINE_IDLE_MINUTES:-10}" # Worker count/memory cap are computed at startup from this machine's actual # resources, not hardcoded — tune these two, not compute_workers() below. -MEM_TARGET_FRACTION="${PIPELINE_MEM_TARGET_FRACTION:-0.55}" -WORKER_MEM_ESTIMATE_MB="${PIPELINE_WORKER_MEM_ESTIMATE_MB:-3200}" +# MEM_TARGET_FRACTION caps TOTAL system memory in use (this pipeline + every +# other process on the box), not just this pipeline's own share — a run that +# starts alone and looks fine can still push the machine to 90%+ once other +# jobs land on the same host, so every check below is against system-wide +# used memory (MemTotal - MemAvailable), never just this pipeline's estimate. +MEM_TARGET_FRACTION="${PIPELINE_MEM_TARGET_FRACTION:-0.70}" +WORKER_MEM_ESTIMATE_MB="${PIPELINE_WORKER_MEM_ESTIMATE_MB:-5500}" STATE="${EXPORT_ROOT}/.pipeline" LOCK="${STATE}/lock" @@ -66,12 +71,28 @@ fi mkdir -p "$STATE" "$CLAIMED" "$FAILED" "$QUARANTINE" "$SCRATCH" "$TMP_OUT" "$LOGS" +# ---- shared helpers ------------------------------------------------------ +log() { echo "[$(date +%H:%M:%S)] [$ROLE] $*"; } + +mem_total_mb() { awk '/MemTotal/{printf "%d", $2/1024}' /proc/meminfo; } + +# System-wide, not this pipeline's own usage: MemAvailable already accounts +# for every process on the box, reclaimable cache included, so MemTotal - +# MemAvailable is what everything combined is actually holding onto right now. +mem_used_mb() { awk '/MemTotal/{t=$2} /MemAvailable/{a=$2} END{printf "%d", (t-a)/1024}' /proc/meminfo; } + +# Budget in MB still free before system-wide usage hits MEM_TARGET_FRACTION +# of total RAM. Negative means already over the target. +mem_budget_mb() { + awk -v total="$(mem_total_mb)" -v used="$(mem_used_mb)" -v f="$MEM_TARGET_FRACTION" \ + 'BEGIN{printf "%d", total*f - used}' +} + # ---- resource sizing (supervisor computes once, workers inherit via env) -- compute_workers() { - local mem_total_kb mem_workers core_workers - mem_total_kb="$(awk '/MemTotal/{print $2}' /proc/meminfo)" - mem_workers="$(awk -v kb="$mem_total_kb" -v f="$MEM_TARGET_FRACTION" -v est="$WORKER_MEM_ESTIMATE_MB" \ - 'BEGIN{printf "%d", (kb/1024*f)/est}')" # awk not $(( )): bash treats leading-zero numbers as octal + local mem_workers core_workers + mem_workers="$(awk -v budget="$(mem_budget_mb)" -v est="$WORKER_MEM_ESTIMATE_MB" \ + 'BEGIN{w=budget/est; printf "%d", (w<0)?0:w}')" # awk not $(( )): bash treats leading-zero numbers as octal core_workers=$(( $(nproc) * 8 / 10 )) # leave headroom, don't claim every core local workers=$mem_workers [[ $core_workers -lt $workers ]] && workers=$core_workers @@ -79,21 +100,14 @@ compute_workers() { echo "$workers" } -mem_available_floor_mb() { - awk -v est="$WORKER_MEM_ESTIMATE_MB" 'BEGIN{printf "%d", est * 1.2}' -} - -# ---- shared helpers ------------------------------------------------------ -log() { echo "[$(date +%H:%M:%S)] [$ROLE] $*"; } - -mem_available_mb() { - awk '/MemAvailable/{printf "%d", $2/1024}' /proc/meminfo -} - +# Blocks a worker from claiming its next batch while starting one more would +# push system-wide memory (this pipeline + anything else running) past +# MEM_TARGET_FRACTION of total RAM. Re-evaluated on every claim, not just at +# startup, so load from unrelated processes throttles new work immediately +# instead of only being noticed once memory is nearly exhausted. wait_for_memory() { - local floor; floor="$(mem_available_floor_mb)" - while [[ "$(mem_available_mb)" -lt "$floor" ]]; do - log "MemAvailable below ${floor}MB, waiting for headroom before claiming more work" + while [[ "$(mem_budget_mb)" -lt "$WORKER_MEM_ESTIMATE_MB" ]]; do + log "system memory usage at $(mem_used_mb)MB/$(mem_total_mb)MB (target ${MEM_TARGET_FRACTION} of total), waiting for headroom before claiming more work" sleep 15 done } @@ -350,8 +364,7 @@ role_supervisor() { echo "$$" > "$PGID_FILE" # setsid above made pid == pgid - WORKERS="$(compute_workers)" - log "sizing: $(nproc) cores, $(awk '/MemTotal/{printf "%.1fGB", $2/1024/1024}' /proc/meminfo) RAM, MEM_TARGET_FRACTION=$MEM_TARGET_FRACTION, WORKER_MEM_ESTIMATE_MB=$WORKER_MEM_ESTIMATE_MB -> WORKERS=$WORKERS" + log "sizing: $(nproc) cores, $(awk '/MemTotal/{printf "%.1fGB", $2/1024/1024}' /proc/meminfo) RAM, MEM_TARGET_FRACTION=$MEM_TARGET_FRACTION, WORKER_MEM_ESTIMATE_MB=$WORKER_MEM_ESTIMATE_MB -> starting at $(compute_workers) workers, will scale with available memory" trap 'log "shutting down"; kill -- -$$ 2>/dev/null || true' INT TERM @@ -359,35 +372,74 @@ role_supervisor() { ( "$0" acquire --mode "$MODE" 2>&1 | sed -u 's/^/[acquire] /' | tee -a "$LOGS/acquire.log" ) & local acquire_pid=$! + # Re-evaluated on every call (not sized once at startup) so the worker count + # tracks memory headroom as it opens up (e.g. other batches committing, + # unrelated processes exiting) instead of being stuck at whatever the launch + # moment happened to allow. Never kills anything to scale down — each + # worker's own wait_for_memory already throttles that side. A fresh index + # per spawn also means a worker that died (OOM, crash) just gets replaced + # here on the next poll rather than needing a full pipeline restart. local -a worker_pids=() - spawn_workers() { - worker_pids=() - local i - for ((i = 0; i < WORKERS; i++)); do - ( "$0" worker --mode "$MODE" "$i" 2>&1 | sed -u "s/^/[w${i}] /" | tee -a "$LOGS/w${i}.log" ) & + local next_worker_idx=0 + top_up_workers() { + local need remaining i + # compute_workers() reads live system-wide usage, which already includes + # every currently-running worker's footprint — so its result IS "how many + # more fit right now", not a total to reconcile against the current + # count. (At startup, with zero workers running, that's the same number + # either way, which is why the original single-shot call worked.) + need="$(compute_workers)" + [[ $need -le 0 ]] && return + # compute_workers() only knows about memory, not remaining work. During + # active collection that's fine — is_candidate's freshness gate means + # list_candidates can read 0 for a moment even with plenty of work still + # to come, so capping on it here would block the very first spawn. + # Once collection is done, though, "no more will ever appear" is a safe + # assumption — cap there so the drain phase can't spawn idle workers + # faster than the ones already idling out finish their exit countdown, + # which would never converge. + if [[ -f "$COLLECTION_DONE" ]]; then + # list_candidates' own exit status is nonzero once nothing is left (its + # last executed statement is a failing `is_candidate && echo` inside a + # for loop) — exactly the case here. A plain assignment doesn't get the + # if/while exemption a bare `[[ ... ]]` condition would, so under + # set -e this kills the whole supervisor unless neutralized. + remaining="$(rebuild_done; list_candidates | wc -l)" || true + [[ $need -gt $remaining ]] && need=$remaining + [[ $need -le 0 ]] && return + fi + [[ ${#worker_pids[@]} -gt 0 ]] && log "memory headroom available ($(mem_used_mb)MB/$(mem_total_mb)MB used): adding $need worker(s) to the ${#worker_pids[@]} running" + for ((i = 0; i < need; i++)); do + ( "$0" worker --mode "$MODE" "$next_worker_idx" 2>&1 | sed -u "s/^/[w${next_worker_idx}] /" | tee -a "$LOGS/w${next_worker_idx}.log" ) & worker_pids+=($!) + next_worker_idx=$((next_worker_idx + 1)) done } - spawn_workers + top_up_workers wait "$acquire_pid" || true - wait "${worker_pids[@]}" || true - # Drain: relaunch until 2 consecutive passes find nothing left (see - # AGENT.md) — a single clean scan could be a transient glitch, and trusting - # it alone here risks merging before everything's actually exported, not - # just losing a worker the way role_worker's equivalent check would. + # Monitor: prune workers that exited (normal drain-out or a crash) and top + # back up while there's still work — same idle-detection reasoning as + # role_worker's empty-scan check (AGENT.md): 2 consecutive dry passes with + # zero live workers and zero candidates before treating the run as done, + # since a single clean scan could be a transient glitch. local drained_scans=0 while [[ $drained_scans -lt 2 ]]; do + local -a alive=() pid + for pid in "${worker_pids[@]}"; do + kill -0 "$pid" 2>/dev/null && alive+=("$pid") + done + worker_pids=("${alive[@]}") + rebuild_done - if [[ -z "$(list_candidates | head -1)" ]] && [[ -z "$(ls -A "$CLAIMED" 2>/dev/null)" ]]; then + if [[ ${#worker_pids[@]} -eq 0 ]] && [[ -z "$(list_candidates | head -1)" ]] && [[ -z "$(ls -A "$CLAIMED" 2>/dev/null)" ]]; then drained_scans=$((drained_scans + 1)) sleep 5 continue fi drained_scans=0 - log "drain pass found leftover work, relaunching workers" - spawn_workers - wait "${worker_pids[@]}" || true + top_up_workers + sleep 60 done local -a batch_dirs=("$EXPORT_ROOT"/batch_*) From 309e3e335dfc22eb32de6f010039b0de7989eb83 Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Tue, 11 Aug 2026 09:13:24 +0200 Subject: [PATCH 06/12] fix: keep worker pool sized against live backlog throughout collection top_up_workers was only ever called once at the start of the collection phase (the supervisor then blocked on `wait "$acquire_pid"` until collection fully finished), so the pool was permanently stuck at whatever system memory allowed at that single instant. Combined with the backlog cap only applying post-collection, a run could spawn its full memory-supported worker count immediately and leave most of it idle for the rest of a long collection if candidates matured slower than the pool could consume them. Now top_up_workers polls every 60s throughout collection too, and caps new spawns against the live unclaimed backlog unconditionally (safe now that a transient empty read just delays one poll instead of blocking forever). Co-Authored-By: Claude Sonnet 5 --- .../worker-pool-static-during-collection.md | 83 +++++++++++++++++++ tools/AGENT.md | 10 +++ tools/pipeline.sh | 47 +++++++---- 3 files changed, 122 insertions(+), 18 deletions(-) create mode 100644 docs/investigations/worker-pool-static-during-collection.md diff --git a/docs/investigations/worker-pool-static-during-collection.md b/docs/investigations/worker-pool-static-during-collection.md new file mode 100644 index 0000000..fd3b603 --- /dev/null +++ b/docs/investigations/worker-pool-static-during-collection.md @@ -0,0 +1,83 @@ +# Symptom: many workers running, but most idle for the whole collection phase + +## Observed + +During a long-running `tools/pipeline.sh` run, `ps aux | grep nova-data-cli` +showed a large number of workers up, but only a handful actively logging +`[worker] committed batch_...` lines — the rest sat silent, doing nothing, +for most of the run. + +## Why it happens + +`role_supervisor` calls `top_up_workers` exactly once, then blocks: + +```bash +top_up_workers +wait "$acquire_pid" || true +``` + +`top_up_workers` sizes the pool from `compute_workers()`, which is a pure +function of **system memory/cores at the instant it's called** — it has no +idea how many recordings are actually claimable yet. At the moment collection +starts, `list_candidates` may already show a nontrivial backlog (e.g. a +`--mode remote` run resuming against a host with data already sitting there), +so this first call can spawn the machine's *entire* memory-supported worker +count immediately — before there's any evidence that candidates will keep +arriving at a matching rate. + +`wait "$acquire_pid"` then blocks the supervisor's own control flow until the +**entire acquisition phase finishes** (which can be hours, for a live +collection). No further `top_up_workers` call happens until then — the +drain-loop's periodic top-up (every 60s) only starts *after* `COLLECTION_DONE` +is set. So whatever pool size the very first call happened to produce is what +runs, unmonitored, for the whole collection window, regardless of how the +real backlog behaves afterward (e.g. if it trickles in far slower than +`workers × CHUNK` can consume, or if more memory frees up later as batches +commit). + +The code comment directly above `top_up_workers` claims "re-evaluated on +every call... so the worker count tracks memory headroom as it opens up" — +true of the *function*, but that guarantee only holds once collection is +already done; during collection, the function is simply never called again +to exercise it. + +## Effect + +Most workers spawned at that first call end up polling `list_candidates` +every 30s and finding nothing (no log line on a non-final empty scan — see +`role_worker`'s `empty_scans` loop — so this is invisible unless you're +watching `ps aux`), while a handful of workers that won the claim race keep +grinding through their batches. Throughput looks fine at a glance (batches do +keep committing), but most of the machine's provisioned worker capacity sits +unused for the run. + +## Fix + +Two changes, both in `tools/pipeline.sh` `role_supervisor`/`top_up_workers`: + +1. Keep calling `top_up_workers` on the same 60s cadence during acquisition + too, not just after `COLLECTION_DONE`: + + ```bash + top_up_workers + while kill -0 "$acquire_pid" 2>/dev/null; do + sleep 60 + top_up_workers + done + wait "$acquire_pid" || true + ``` + +2. `top_up_workers` now caps the new-worker count against the actual + unclaimed backlog (`list_candidates | wc -l`) on every call, not only once + `COLLECTION_DONE` is set. This cap already existed for the post-collection + drain phase; it was deliberately *not* applied during collection because a + single `list_candidates` read coming back empty (e.g. `is_candidate`'s + 60s-untouched freshness gate momentarily reading zero) would otherwise + permanently block the very first spawn under the old one-shot-call + structure. Change (1) removes that risk: a transient zero reading now just + costs one 60s poll, not the rest of the run, so the cap can safely apply + throughout. + +Together these mean the pool actually scales with the live backlog and +memory headroom as collection proceeds, instead of being frozen at whatever +the single startup snapshot allowed. diff --git a/tools/AGENT.md b/tools/AGENT.md index 9b72e88..184db84 100644 --- a/tools/AGENT.md +++ b/tools/AGENT.md @@ -152,6 +152,16 @@ keeps this from monopolizing CPU/disk even without an active watchdog for those resources — it tells the kernel to prefer any other process, so the pipeline only consumes spare capacity. +`top_up_workers` is polled every 60s for the pipeline's *entire* life — +during acquisition as well as the post-`COLLECTION_DONE` drain loop, not just +the latter — and on every call caps how many new workers it spawns against +the live unclaimed backlog (`list_candidates`), not just available memory. +Without both of these, the pool gets sized once, from whatever memory allowed +at the instant collection started, and never adjusts again for the rest of a +run that can last hours: if candidates become claimable slower than +`workers × CHUNK` can consume, most of that pool just sits idle the whole +time (see `docs/investigations/worker-pool-static-during-collection.md`). + ## "Nothing left" needs confirmation, not a single scan A worker exits when it finds no candidates *and* `COLLECTION_DONE` exists; the diff --git a/tools/pipeline.sh b/tools/pipeline.sh index 1b98dfd..bad3a5b 100755 --- a/tools/pipeline.sh +++ b/tools/pipeline.sh @@ -390,24 +390,25 @@ role_supervisor() { # either way, which is why the original single-shot call worked.) need="$(compute_workers)" [[ $need -le 0 ]] && return - # compute_workers() only knows about memory, not remaining work. During - # active collection that's fine — is_candidate's freshness gate means - # list_candidates can read 0 for a moment even with plenty of work still - # to come, so capping on it here would block the very first spawn. - # Once collection is done, though, "no more will ever appear" is a safe - # assumption — cap there so the drain phase can't spawn idle workers - # faster than the ones already idling out finish their exit countdown, - # which would never converge. - if [[ -f "$COLLECTION_DONE" ]]; then - # list_candidates' own exit status is nonzero once nothing is left (its - # last executed statement is a failing `is_candidate && echo` inside a - # for loop) — exactly the case here. A plain assignment doesn't get the - # if/while exemption a bare `[[ ... ]]` condition would, so under - # set -e this kills the whole supervisor unless neutralized. - remaining="$(rebuild_done; list_candidates | wc -l)" || true - [[ $need -gt $remaining ]] && need=$remaining - [[ $need -le 0 ]] && return - fi + # compute_workers() only knows about memory, not remaining work — cap the + # new-worker count against the actual unclaimed backlog too, or a run + # whose candidates trickle in slower than memory allows workers ends up + # with most of them spawned up front and idling for the rest of + # collection (see docs/investigations/worker-pool-static-during-collection.md). + # Safe to do unconditionally (not just once collection is done): this + # function is now polled every 60s throughout collection too (see the + # loop around its first call below), so a transient zero-candidate + # reading (e.g. is_candidate's freshness gate momentarily reading empty) + # just delays the next top-up by one poll instead of blocking it forever. + # + # list_candidates' own exit status is nonzero once nothing is left (its + # last executed statement is a failing `is_candidate && echo` inside a + # for loop) — a plain assignment doesn't get the if/while exemption a + # bare `[[ ... ]]` condition would, so under set -e this kills the whole + # supervisor unless neutralized. + remaining="$(rebuild_done; list_candidates | wc -l)" || true + [[ $need -gt $remaining ]] && need=$remaining + [[ $need -le 0 ]] && return [[ ${#worker_pids[@]} -gt 0 ]] && log "memory headroom available ($(mem_used_mb)MB/$(mem_total_mb)MB used): adding $need worker(s) to the ${#worker_pids[@]} running" for ((i = 0; i < need; i++)); do ( "$0" worker --mode "$MODE" "$next_worker_idx" 2>&1 | sed -u "s/^/[w${next_worker_idx}] /" | tee -a "$LOGS/w${next_worker_idx}.log" ) & @@ -415,7 +416,17 @@ role_supervisor() { next_worker_idx=$((next_worker_idx + 1)) done } + # Keep topping up while acquisition is still running, not just once at + # startup — otherwise the pool is permanently stuck at whatever memory + # happened to allow the instant collection began (often oversized relative + # to how fast candidates actually become claimable), for the entire + # collection phase, however long that is. Same poll cadence as the drain + # loop below. top_up_workers + while kill -0 "$acquire_pid" 2>/dev/null; do + sleep 60 + top_up_workers + done wait "$acquire_pid" || true # Monitor: prune workers that exited (normal drain-out or a crash) and top From 3d38751e31d536a744f8b16421c819b1dbef3510 Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Tue, 11 Aug 2026 11:29:35 +0200 Subject: [PATCH 07/12] docs: design and plan for per-episode task instructions + generalized metadata Design spec (verified against the installed lerobot package, LeRobot policy configs, and real community datasets like DROID and AgiBot World 2026) plus the resulting task-by-task implementation plan for two optional exporter features: per-episode language instructions (config.task_field, sourced from meta.json, falling back to the existing dataset-wide task_description) and a fix generalizing episode_metadata to any JSON scalar type instead of float-only with a wrong 0.0 default for missing fields. Co-Authored-By: Claude Sonnet 5 --- ...026-08-11-per-episode-task-and-metadata.md | 721 ++++++++++++++++++ ...11-per-episode-task-and-metadata-design.md | 174 +++++ 2 files changed, 895 insertions(+) create mode 100644 docs/superpowers/plans/2026-08-11-per-episode-task-and-metadata.md create mode 100644 docs/superpowers/specs/2026-08-11-per-episode-task-and-metadata-design.md diff --git a/docs/superpowers/plans/2026-08-11-per-episode-task-and-metadata.md b/docs/superpowers/plans/2026-08-11-per-episode-task-and-metadata.md new file mode 100644 index 0000000..87e4691 --- /dev/null +++ b/docs/superpowers/plans/2026-08-11-per-episode-task-and-metadata.md @@ -0,0 +1,721 @@ +# Per-episode task instruction + generalized episode metadata Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Let the exporter vary the LeRobot `task` string per episode (read from each recording's `meta.json`, e.g. `"task"`) instead of one fixed dataset-wide string, and fix the existing `episode_metadata` mechanism to handle any JSON scalar type (not just floats) — both fully optional, defaulting to today's exact behavior when unset. + +**Architecture:** Reuse the existing `meta.json`-reading path in `exporter.py` (`_load_episode_metadata`) for both features — extend it to also resolve a per-episode task string, and fix its value typing/missing-field default along the way. Thread the resolved value through `Episode.task` (new field, `episode_sampler.py`) into `LeRobotHead._sample_to_frame` (`heads/lerobot.py`), which already writes `frame["task"]` per sample — it just currently always reads `config.task_description` there instead of a per-episode value. + +**Tech Stack:** Python 3.13, pydantic (`ExportConfig`), pytest. + +## Global Constraints + +- `task_field` defaults to `None`; `episode_metadata` defaults to `[]`. With both unset, export output must be byte-for-byte identical to current behavior — every task must preserve this. +- Per-episode, not per-frame: the resolved task string is constant across all frames of one episode (matches how LeRobot's `task`/`tasks.parquet`/`task_index` mechanism actually works — see `docs/superpowers/specs/2026-08-11-per-episode-task-and-metadata-design.md`). +- `episode_metadata` values are `Any` JSON scalar (str/float/int/bool) — no schema/type validation across episodes. +- Missing `meta.json` field (or no local `meta.json` at all, e.g. `catalog_url` export): warn and fall back — `task_field` falls back to `config.task_description`; `episode_metadata` fields fall back to `None` (not `0.0`). Never a hard failure. +- Out of scope: `heads/groot.py` (unaffected — it already gets per-episode tasks for free via LeRobot's `task_index`), and LeRobot's `language_persistent`/`language_events` schema (a separate, unrelated annotation mechanism). + +--- + +## File Structure + +- Modify: `src/nova_export/export/config.py` — add `ExportConfig.task_field`. +- Modify: `src/nova_export/export/episode_sampler.py` — add `Episode.task`; widen `Episode.extra_metadata`'s type. +- Modify: `src/nova_export/export/exporter.py` — generalize `_load_episode_metadata` (typing + task_field support); resolve `episode.task`/`episode.extra_metadata` in `export_recordings`. +- Modify: `src/nova_export/export/heads/lerobot.py` — `_sample_to_frame`/`write_episode` use the resolved per-episode task instead of the config constant. +- Modify: `tests/test_export.py` — new/updated tests for all of the above. +- Modify: `docs/export-guide.md` — document `task_field`; correct `episode_metadata`'s description. + +--- + +### Task 1: Add `ExportConfig.task_field` + +**Files:** +- Modify: `src/nova_export/export/config.py:156-176` +- Test: `tests/test_export.py` (new test in a `TestExportConfig`-style block, or alongside existing config-via-`ExportConfig(...)` tests — there's no dedicated `TestExportConfig` class yet; add one near the top of the file, after the imports/helpers, before `TestFrameCache`) + +**Interfaces:** +- Produces: `ExportConfig.task_field: str | None` (default `None`) — consumed by Task 3. + +- [ ] **Step 1: Write the failing test** + +Add to `tests/test_export.py` (a new class, placed after the helper functions around line 172 and before `class TestFrameCache:`): + +```python +class TestExportConfigTaskField: + """Tests for ExportConfig.task_field.""" + + def test_task_field_defaults_to_none(self): + config = ExportConfig(fps=15) + assert config.task_field is None + + def test_task_field_can_be_set(self): + config = ExportConfig(fps=15, task_field="task") + assert config.task_field == "task" +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py::TestExportConfigTaskField -v` +Expected: FAIL — `ExportConfig` has no field `task_field` (pydantic raises on the second test since it's an unknown kwarg... actually pydantic's default `extra` behavior is to raise `ValidationError` for unknown fields, or the first test fails with `AttributeError: 'ExportConfig' object has no attribute 'task_field'`). + +- [ ] **Step 3: Add the field** + +In `src/nova_export/export/config.py`, right after the existing `task_description` field (currently lines 156-159): + +```python + task_description: str = Field( + default="task", + description="Task label written to the dataset", + ) + + task_field: str | None = Field( + default=None, + description=( + "meta.json field name (e.g. 'task') holding this episode's " + "natural-language task instruction. When set, each episode's " + "LeRobot 'task' is read from its own meta.json instead of the " + "fixed task_description, mirroring how LeRobot's task/task_index " + "mechanism is meant to vary per episode. Falls back to " + "task_description when the field is missing for a given episode, " + "or when no local meta.json is available (e.g. exporting from " + "catalog_url). Requires local rrd_paths export, like " + "episode_metadata." + ), + ) +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py::TestExportConfigTaskField -v` +Expected: PASS + +- [ ] **Step 5: Commit** + +```bash +git add src/nova_export/export/config.py tests/test_export.py +git commit -m "feat: add optional ExportConfig.task_field for per-episode task instructions" +``` + +--- + +### Task 2: Generalize `episode_metadata` beyond floats + +**Files:** +- Modify: `src/nova_export/export/episode_sampler.py:66` +- Modify: `src/nova_export/export/exporter.py:203-231` (`_load_episode_metadata`) and `:491-503` (the `episode.extra_metadata` assignment inside `export_recordings`) +- Test: `tests/test_export.py` + +**Interfaces:** +- Consumes: nothing new from Task 1. +- Produces: `_resolve_extra_metadata(found: dict[str, Any], fields: list[str]) -> dict[str, Any]` — a new module-level function in `exporter.py`, consumed by `export_recordings` in this same task and left untouched by Task 3. `Episode.extra_metadata: dict[str, Any] | None` (was `dict[str, float] | None`). +- Note for the implementer: `export_recordings`'s metadata block currently builds `episode.extra_metadata` inline (`{f: found.get(f, 0.0) for f in config.episode_metadata}`) — that inline dict comprehension is *not* independently unit-testable (it's buried inside a large loop over live Rerun segments), which is exactly why the bug this task fixes (`0.0` instead of `None` for a missing field) has no direct test today. This task extracts that expression into `_resolve_extra_metadata`, a pure function, specifically so the fix has a real test that fails before the fix and passes after — not a test that merely documents the intended expression. + +- [ ] **Step 1: Write the failing test** + +Add to `tests/test_export.py`, inside a new `TestResolveExtraMetadata` class (place it right before `class TestSourceValidation:`, since it tests another `exporter.py` private helper the same way that class does): + +```python +class TestResolveExtraMetadata: + """Tests for exporter._resolve_extra_metadata.""" + + def test_present_fields_pass_through_any_type(self): + from nova_export.export.exporter import _resolve_extra_metadata + + found = {"cube_x_mm": -324.15, "cube_color": "purple"} + + result = _resolve_extra_metadata(found, ["cube_x_mm", "cube_color"]) + + assert result == {"cube_x_mm": -324.15, "cube_color": "purple"} + + def test_missing_field_defaults_to_none_not_zero(self): + from nova_export.export.exporter import _resolve_extra_metadata + + found = {"cube_x_mm": 1.5} # cube_color absent + + result = _resolve_extra_metadata(found, ["cube_x_mm", "cube_color"]) + + assert result == {"cube_x_mm": 1.5, "cube_color": None} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py::TestResolveExtraMetadata -v` +Expected: FAIL with `ImportError: cannot import name '_resolve_extra_metadata'` — the function doesn't exist yet. + +- [ ] **Step 3: Add `_resolve_extra_metadata` and fix the typing** + +In `src/nova_export/export/episode_sampler.py`, change line 66: + +```python + extra_metadata: dict[str, float] | None = None +``` + +to: + +```python + extra_metadata: dict[str, Any] | None = None +``` + +(`Any` is already imported at the top of this file.) + +In `src/nova_export/export/exporter.py`, add `from typing import Any` near the top (this file currently has no `typing` import at all — add it next to the other stdlib imports), then add this new function directly above `_load_episode_metadata` (currently at line 203): + +```python +def _resolve_extra_metadata(found: dict[str, Any], fields: list[str]) -> dict[str, Any]: + """Build one episode's extra_metadata dict from its meta.json values. + + A field missing from `found` becomes None, not 0.0 — 0.0 was only + correct by accident for numeric fields and actively wrong for a field + like cube_color ("purple"). + """ + return {f: found.get(f) for f in fields} +``` + +Also update `_load_episode_metadata`'s return-type annotation only (its body is already type-agnostic — the bug was purely in the caller): + +```python +def _load_episode_metadata( + rrd_paths: list[Path] | None, fields: list[str] +) -> dict[str, dict[str, Any]]: + """Read episode_metadata fields from each recording's sibling meta.json. + + Keyed by segment_id, which for local rrd_paths exports is exactly the + recording's directory name (//recording.rrd) — + the same recording_id the collector assigns and rerun uses as the + segment ID, so no separate ID plumbing is needed. Values are whatever + JSON scalar type meta.json holds (str/float/int/bool) — not float-only. + """ + if not fields: + return {} + if not rrd_paths: + logger.warning( + "episode_metadata {} configured but exporting from catalog_url " + "(no local meta.json available) — skipping", + fields, + ) + return {} + + result: dict[str, dict[str, Any]] = {} + for rrd_path in rrd_paths: + meta_path = rrd_path.parent / "meta.json" + if not meta_path.is_file(): + continue + meta = json.loads(meta_path.read_text()) + segment_id = rrd_path.parent.name + result[segment_id] = {f: meta[f] for f in fields if f in meta} + return result +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py::TestResolveExtraMetadata -v` +Expected: PASS + +- [ ] **Step 5: Wire `_resolve_extra_metadata` into `export_recordings`** + +In `src/nova_export/export/exporter.py`, change the caller (currently lines 491-503) — only the last two lines of this block change (the log message text, and the final assignment now calling the new function): + +```python + if config.episode_metadata: + found = episode_metadata_by_segment.get(segment_id, {}) + missing = [f for f in config.episode_metadata if f not in found] + if missing: + logger.warning( + "Episode {} ({}): meta.json missing {} — filled with null", + episode_id, + segment_id[:8], + missing, + ) + episode.extra_metadata = _resolve_extra_metadata( + found, config.episode_metadata + ) +``` + +- [ ] **Step 6: Run the full test file to verify nothing regressed** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py -v -k "not RealIntegration"` +Expected: PASS (the `-k "not RealIntegration"` skips the slow real-`.rrd` integration test, which needs real recording fixtures and isn't affected by this change anyway) + +- [ ] **Step 7: Commit** + +```bash +git add src/nova_export/export/episode_sampler.py src/nova_export/export/exporter.py tests/test_export.py +git commit -m "fix: generalize episode_metadata to any JSON scalar, default missing to null" +``` + +--- + +### Task 3: Resolve per-episode `task` in `exporter.py` + +**Files:** +- Modify: `src/nova_export/export/episode_sampler.py:60-66` (`Episode` dataclass) +- Modify: `src/nova_export/export/exporter.py:203-231` (`_load_episode_metadata`, extended) and `:395-397` + the loop body around `:491-503` +- Test: `tests/test_export.py` + +**Interfaces:** +- Consumes: `ExportConfig.task_field` (Task 1), `_resolve_extra_metadata` and `Episode.extra_metadata: dict[str, Any] | None` (Task 2). +- Produces: `Episode.task: str | None` (new field, default `None`) — consumed by Task 4. `_resolve_task(found: dict[str, Any], task_field: str | None, task_description: str) -> str` — a new module-level function in `exporter.py`, mirroring `_resolve_extra_metadata`'s shape, for the same reason (a directly-testable pure function instead of inline logic buried in the export loop). `_load_episode_metadata(rrd_paths, fields, task_field=None) -> dict[str, dict[str, Any]]` — the per-segment dict now also carries the task_field's raw value under its own key (same dict, no new return shape) when `task_field` is passed. + +- [ ] **Step 1: Write the failing test** + +Add to `tests/test_export.py`, in a new `TestResolveTask` class placed right after `TestResolveExtraMetadata`: + +```python +class TestResolveTask: + """Tests for exporter._resolve_task.""" + + def test_task_field_unset_uses_task_description(self): + from nova_export.export.exporter import _resolve_task + + result = _resolve_task({}, None, "fallback_task") + + assert result == "fallback_task" + + def test_task_field_present_used_verbatim(self): + from nova_export.export.exporter import _resolve_task + + found = {"task": "Pick the purple cube up."} + + result = _resolve_task(found, "task", "fallback_task") + + assert result == "Pick the purple cube up." + + def test_task_field_missing_from_found_falls_back(self): + from nova_export.export.exporter import _resolve_task + + found = {"other_field": 1} # no "task" key + + result = _resolve_task(found, "task", "fallback_task") + + assert result == "fallback_task" +``` + +Also add these two cases to `TestLoadEpisodeMetadata` (rename `TestLoadEpisodeMetadata` if it doesn't already exist under that name — it was not introduced in Task 2, since Task 2's tests live in `TestResolveExtraMetadata` instead; create `TestLoadEpisodeMetadata` fresh here, placed right after `TestResolveTask`): + +```python +class TestLoadEpisodeMetadata: + """Tests for exporter._load_episode_metadata.""" + + def test_string_field_round_trips(self, tmp_path): + from nova_export.export.exporter import _load_episode_metadata + + rec_dir = tmp_path / "04cb4f25d3ef" + rec_dir.mkdir() + (rec_dir / "meta.json").write_text( + '{"cube_color": "purple", "cube_x_mm": -324.15}' + ) + rrd_path = rec_dir / "recording.rrd" + rrd_path.touch() + + result = _load_episode_metadata([rrd_path], ["cube_color", "cube_x_mm"]) + + assert result["04cb4f25d3ef"]["cube_color"] == "purple" + assert result["04cb4f25d3ef"]["cube_x_mm"] == -324.15 + + def test_missing_field_simply_absent_from_result(self, tmp_path): + from nova_export.export.exporter import _load_episode_metadata + + rec_dir = tmp_path / "rec01" + rec_dir.mkdir() + (rec_dir / "meta.json").write_text('{"cube_x_mm": 1.5}') + rrd_path = rec_dir / "recording.rrd" + rrd_path.touch() + + result = _load_episode_metadata([rrd_path], ["cube_x_mm", "cube_color"]) + + assert result["rec01"] == {"cube_x_mm": 1.5} + assert "cube_color" not in result["rec01"] + + def test_task_field_value_included_in_result(self, tmp_path): + from nova_export.export.exporter import _load_episode_metadata + + rec_dir = tmp_path / "04cb4f25d3ef" + rec_dir.mkdir() + (rec_dir / "meta.json").write_text( + '{"task": "Pick the purple cube up.", "cube_x_mm": 1.0}' + ) + rrd_path = rec_dir / "recording.rrd" + rrd_path.touch() + + result = _load_episode_metadata([rrd_path], ["cube_x_mm"], task_field="task") + + assert result["04cb4f25d3ef"]["task"] == "Pick the purple cube up." + assert result["04cb4f25d3ef"]["cube_x_mm"] == 1.0 + + def test_task_field_none_does_not_add_task_key(self, tmp_path): + from nova_export.export.exporter import _load_episode_metadata + + rec_dir = tmp_path / "rec01" + rec_dir.mkdir() + (rec_dir / "meta.json").write_text('{"task": "unused", "cube_x_mm": 1.0}') + rrd_path = rec_dir / "recording.rrd" + rrd_path.touch() + + result = _load_episode_metadata([rrd_path], ["cube_x_mm"]) # task_field omitted + + assert result["rec01"] == {"cube_x_mm": 1.0} +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py::TestResolveTask tests/test_export.py::TestLoadEpisodeMetadata -v` +Expected: FAIL — `TestResolveTask` fails with `ImportError: cannot import name '_resolve_task'`; `TestLoadEpisodeMetadata::test_task_field_value_included_in_result` fails with `TypeError: _load_episode_metadata() got an unexpected keyword argument 'task_field'`. The other two `TestLoadEpisodeMetadata` cases (`test_string_field_round_trips`, `test_missing_field_simply_absent_from_result`) already pass against Task 2's code — that's expected, they're regression coverage for behavior Task 2 already delivered, not new-in-this-task assertions. + +- [ ] **Step 3: Add `_resolve_task` and extend `_load_episode_metadata`** + +In `src/nova_export/export/exporter.py`, add this function directly below `_resolve_extra_metadata`: + +```python +def _resolve_task( + found: dict[str, Any], task_field: str | None, task_description: str +) -> str: + """Resolve one episode's task string. + + task_field's value from meta.json when set and present; task_description + otherwise (unset task_field, or the field missing from this episode's + meta.json) — the same fallback either way, so a per-recording gap in + metadata degrades to today's dataset-wide constant rather than failing. + """ + if not task_field: + return task_description + value = found.get(task_field) + return task_description if value is None else str(value) +``` + +Then replace `_load_episode_metadata` (as left by Task 2) with: + +```python +def _load_episode_metadata( + rrd_paths: list[Path] | None, + fields: list[str], + task_field: str | None = None, +) -> dict[str, dict[str, Any]]: + """Read episode_metadata fields (and optionally task_field) from each + recording's sibling meta.json. + + Keyed by segment_id, which for local rrd_paths exports is exactly the + recording's directory name (//recording.rrd) — + the same recording_id the collector assigns and rerun uses as the + segment ID, so no separate ID plumbing is needed. Values are whatever + JSON scalar type meta.json holds (str/float/int/bool) — not float-only. + + When task_field is set, its value (if present) is included in the + per-segment dict under its own key, alongside the requested + episode_metadata fields — one meta.json read serves both. + """ + if not fields and not task_field: + return {} + if not rrd_paths: + logger.warning( + "episode_metadata {} / task_field {!r} configured but exporting " + "from catalog_url (no local meta.json available) — skipping", + fields, + task_field, + ) + return {} + + result: dict[str, dict[str, Any]] = {} + for rrd_path in rrd_paths: + meta_path = rrd_path.parent / "meta.json" + if not meta_path.is_file(): + continue + meta = json.loads(meta_path.read_text()) + segment_id = rrd_path.parent.name + entry: dict[str, Any] = {f: meta[f] for f in fields if f in meta} + if task_field and task_field in meta: + entry[task_field] = meta[task_field] + result[segment_id] = entry + return result +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py::TestResolveTask tests/test_export.py::TestLoadEpisodeMetadata -v` +Expected: PASS (all cases) + +- [ ] **Step 5: Add `Episode.task` and wire resolution into `export_recordings`** + +In `src/nova_export/export/episode_sampler.py`, change the `Episode` dataclass (currently lines 60-66) to: + +```python +class Episode: + """A complete episode with metadata and samples.""" + + segment_id: str + episode_index: int + samples: list[Sample] + extra_metadata: dict[str, Any] | None = None + task: str | None = None +``` + +In `src/nova_export/export/exporter.py`, change the call site (currently lines 395-397): + +```python + episode_metadata_by_segment = _load_episode_metadata( + rrd_paths, config.episode_metadata, config.task_field + ) +``` + +And change the loop body's metadata block (as left by Task 2, currently around lines 491-503) to also resolve `episode.task`, right after the `episode.extra_metadata` assignment — `found` moves out of the `if config.episode_metadata:` guard since both blocks below need it: + +```python + found = episode_metadata_by_segment.get(segment_id, {}) + + if config.episode_metadata: + missing = [f for f in config.episode_metadata if f not in found] + if missing: + logger.warning( + "Episode {} ({}): meta.json missing {} — filled with null", + episode_id, + segment_id[:8], + missing, + ) + episode.extra_metadata = _resolve_extra_metadata( + found, config.episode_metadata + ) + + if config.task_field and config.task_field not in found: + logger.warning( + "Episode {} ({}): meta.json missing task_field {!r} " + "— falling back to task_description", + episode_id, + segment_id[:8], + config.task_field, + ) + episode.task = _resolve_task( + found, config.task_field, config.task_description + ) +``` + +- [ ] **Step 6: Run the full test file to verify nothing regressed** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py -v -k "not RealIntegration"` +Expected: PASS + +- [ ] **Step 7: Commit** + +```bash +git add src/nova_export/export/episode_sampler.py src/nova_export/export/exporter.py tests/test_export.py +git commit -m "feat: resolve per-episode task from meta.json via config.task_field" +``` + +--- + +### Task 4: Write the resolved task in `LeRobotHead`, update docs + +**Files:** +- Modify: `src/nova_export/export/heads/lerobot.py:128-171` (`write_episode`) and `:230-255` (`_sample_to_frame`) +- Modify: `docs/export-guide.md` (config reference table) +- Test: `tests/test_export.py` (`TestLeRobotHead`) + +**Interfaces:** +- Consumes: `Episode.task: str | None` (Task 3). +- Produces: nothing new downstream — this is the terminal consumer of `episode.task`. + +- [ ] **Step 1: Write the failing test** + +Add to `TestLeRobotHead` in `tests/test_export.py`, right after `test_write_episode` (currently ending at line 710): + +```python + @patch("lerobot.datasets.lerobot_dataset.LeRobotDataset") + def test_write_episode_uses_per_episode_task(self, mock_dataset_cls): + """When Episode.task is set, every frame's 'task' must use it — + not the dataset-wide config.task_description.""" + mock_dataset = MagicMock() + mock_dataset_cls.create.return_value = mock_dataset + + config = ExportConfig(fps=15, task_description="fallback_task") + + with tempfile.TemporaryDirectory() as tmpdir: + head = LeRobotHead(config, Path(tmpdir) / "output") + head.initialize({"action": {"dtype": "float32", "shape": (7,)}}) + + episode = create_test_episode(num_samples=3) + episode.task = "Pick the purple cube up." + head.write_episode(episode) + + for call in mock_dataset.add_frame.call_args_list: + frame = call.args[0] + assert frame["task"] == "Pick the purple cube up." + + @patch("lerobot.datasets.lerobot_dataset.LeRobotDataset") + def test_write_episode_falls_back_to_task_description(self, mock_dataset_cls): + """When Episode.task is unset (None), fall back to config.task_description — + this is the byte-for-byte-identical-to-today path.""" + mock_dataset = MagicMock() + mock_dataset_cls.create.return_value = mock_dataset + + config = ExportConfig(fps=15, task_description="fallback_task") + + with tempfile.TemporaryDirectory() as tmpdir: + head = LeRobotHead(config, Path(tmpdir) / "output") + head.initialize({"action": {"dtype": "float32", "shape": (7,)}}) + + episode = create_test_episode(num_samples=3) # episode.task defaults to None + + head.write_episode(episode) + + for call in mock_dataset.add_frame.call_args_list: + frame = call.args[0] + assert frame["task"] == "fallback_task" + + @patch("lerobot.datasets.lerobot_dataset.LeRobotDataset") + def test_different_episodes_can_have_different_tasks(self, mock_dataset_cls): + """Two episodes with distinct task strings both write their own value — + proves this is LeRobot's per-episode task mechanism, not a renamed + per-dataset constant.""" + mock_dataset = MagicMock() + mock_dataset_cls.create.return_value = mock_dataset + + config = ExportConfig(fps=15) + + with tempfile.TemporaryDirectory() as tmpdir: + head = LeRobotHead(config, Path(tmpdir) / "output") + head.initialize({"action": {"dtype": "float32", "shape": (7,)}}) + + episode_a = create_test_episode(segment_id="a", num_samples=1) + episode_a.task = "Task A" + episode_b = create_test_episode(segment_id="b", num_samples=1) + episode_b.task = "Task B" + + head.write_episode(episode_a) + head.write_episode(episode_b) + + frame_a = mock_dataset.add_frame.call_args_list[0].args[0] + frame_b = mock_dataset.add_frame.call_args_list[1].args[0] + assert frame_a["task"] == "Task A" + assert frame_b["task"] == "Task B" +``` + +- [ ] **Step 2: Run test to verify it fails** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py::TestLeRobotHead::test_write_episode_uses_per_episode_task -v` +Expected: FAIL — `frame["task"]` is `"fallback_task"` (from `config.task_description`) instead of `"Pick the purple cube up."`, since `_sample_to_frame` doesn't look at `episode.task` yet. + +- [ ] **Step 3: Thread the resolved task through `write_episode`/`_sample_to_frame`** + +In `src/nova_export/export/heads/lerobot.py`, change `write_episode` (currently lines 128-171) — only the body between the logging call and the `try` block changes: + +```python + def write_episode(self, episode: Episode) -> bool: + """Write an episode to the dataset. + + Args: + episode: Episode to write. + + Returns: + True if successfully written. + """ + if self._dataset is None: + raise RuntimeError("Dataset not initialized. Call initialize() first.") + + if not episode.samples: + logger.warning("Skipping empty episode {}", episode.episode_index) + return False + + logger.info( + "Writing episode {}: {} frames, {:.2f}s duration", + episode.episode_index, + episode.num_frames, + episode.duration_s, + ) + + if self.config.episode_metadata and episode.extra_metadata: + # LeRobot assigns its own sequential episode_index (meta.total_episodes) + # when save_episode() runs, which is NOT episode.episode_index (that's + # the exporter's raw segment-loop counter, and diverges as soon as any + # earlier segment is skipped). Key by the index LeRobot is about to use. + self._episode_metadata_by_index[self._dataset.meta.total_episodes] = ( + episode.extra_metadata + ) + + task = episode.task if episode.task is not None else self.config.task_description + + try: + for sample in tqdm(episode.samples, desc="Frames", leave=False): + frame = self._sample_to_frame(sample, task) + self._dataset.add_frame(frame) + + self._dataset.save_episode() + self._update_counts(episode) + return True + + except Exception as e: + logger.error("Error writing episode {}: {}", episode.episode_index, e) + return False +``` + +And `_sample_to_frame` (currently lines 230-255): + +```python + def _sample_to_frame(self, sample: Sample, task: str) -> dict[str, Any]: + """Convert a Sample to a LeRobot frame dict. + + Args: + sample: Sample to convert. + task: Resolved task string for this sample's episode (either + Episode.task, when set, or config.task_description). + + Returns: + Frame dict for LeRobotDataset.add_frame(). + """ + frame: dict[str, Any] = {} + + # Action + frame["action"] = sample.action + + # State + if len(sample.state) > 0: + frame["observation.state"] = sample.state + + # Task + frame["task"] = task + + # Images + for cam_name, img_array in sample.images.items(): + frame[f"observation.images.{cam_name}"] = img_array + + return frame +``` + +- [ ] **Step 4: Run test to verify it passes** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py::TestLeRobotHead -v` +Expected: PASS (all `TestLeRobotHead` tests, including the three new ones and the pre-existing `test_write_episode` — that one still passes since `create_test_episode` leaves `episode.task=None`, falling back to `config.task_description`, exactly as it wrote before this change) + +- [ ] **Step 5: Run the full test file** + +Run: `cd /home/sebi/ws/nova-data-cli && uv run pytest tests/test_export.py -v -k "not RealIntegration"` +Expected: PASS + +- [ ] **Step 6: Update `docs/export-guide.md`** + +In the config reference table, change the `task_description` row's wording slightly and add a `task_field` row right after it (find the row starting with `| \`task_description\`` in the table around line 69): + +```markdown +| `task_description` | string | `"task"` | Fallback task label written to every frame when `task_field` is unset (or its meta.json field is missing for a given episode). | +| `task_field` | string \| null | `null` | meta.json field name (e.g. `"task"`) holding each episode's own natural-language instruction — lets `task` vary per episode instead of being fixed dataset-wide. Falls back to `task_description`. Requires local export (same as `episode_metadata`). | +``` + +Also fix the existing `episode_metadata` row's description (search for `episode_metadata` in the table) to drop any float-specific wording if present, or add a one-line clarification directly below the table: + +```markdown +`episode_metadata` values may be any JSON scalar type (string, number, boolean) — not float-only. A field missing from a given episode's `meta.json` is filled with `null` in that episode's row. +``` + +- [ ] **Step 7: Commit** + +```bash +git add src/nova_export/export/heads/lerobot.py tests/test_export.py docs/export-guide.md +git commit -m "feat: write per-episode task string in LeRobotHead, document task_field" +``` + +--- + +## Self-Review Notes + +- **Spec coverage:** Data flow section (meta.json → episode.task → frame["task"], and meta.json → episode.extra_metadata → parquet columns) — Tasks 2-4. Error handling (missing field/no meta.json → fallback + warning, never a hard failure) — Task 3 Step 5. Testing section's five `_load_episode_metadata`-level cases and the LeRobotHead-level multi-task-string case — Tasks 2-4's test steps. Optionality requirement (both features default to today's exact behavior) — every task's fallback-path test (`test_write_episode_falls_back_to_task_description`, Task 2's unaffected-when-unset defaults). +- **Type consistency checked:** `Episode.extra_metadata: dict[str, Any] | None` (Task 2) and `Episode.task: str | None` (Task 3) match how `exporter.py` sets them (Task 3 Step 5) and how `heads/lerobot.py` reads them (Task 4 Step 3). `_load_episode_metadata`'s signature gains `task_field` in Task 3 without breaking Task 2's call sites (default `None` keeps the two-arg call from Task 2's own tests working). +- **No placeholders:** every step has literal code, not a description of code. diff --git a/docs/superpowers/specs/2026-08-11-per-episode-task-and-metadata-design.md b/docs/superpowers/specs/2026-08-11-per-episode-task-and-metadata-design.md new file mode 100644 index 0000000..8c032e8 --- /dev/null +++ b/docs/superpowers/specs/2026-08-11-per-episode-task-and-metadata-design.md @@ -0,0 +1,174 @@ +# Per-episode language instructions + generalized episode metadata + +## Problem + +The exporter needs to support VLA training datasets, where each episode carries +its own natural-language task instruction (e.g. "Pick the purple cube up from +the left of the table and place it onto the yellow square target position."), +alongside per-episode debug metadata (e.g. where the cube actually started: +`cube_x_mm`, `cube_y_mm`, `cube_z_mm`, `cube_yaw_rad`, `cube_color`). + +Today: + +- `config.task_description` is a single fixed string written to *every* frame + of *every* episode (`heads/lerobot.py:_sample_to_frame`) — there is no way + to vary it per episode. +- `config.episode_metadata` (a list of `meta.json` field names) already reads + per-episode values and writes them as extra columns on + `meta/episodes/*.parquet` — but `exporter.py:_load_episode_metadata` types + values as `dict[str, float]` and defaults a missing field to `0.0`. Real + `meta.json` data includes non-numeric fields (`cube_color: "purple"`), + which this silently mishandles. + +Both mechanisms must stay **optional**: a plain imitation-learning dataset +with no task and no extra metadata must export exactly as it does today — +nothing here is a new requirement. + +## Non-goals + +- No per-frame-varying task (only per-episode, matching how LeRobot's + `task`/`task_index` mechanism actually works). +- No schema/typing for `episode_metadata` fields beyond "whatever JSON scalar + is in `meta.json`" (str / float / int / bool) — no validation that a field + is consistently typed across episodes. +- No change to GR00T's own code — it already gets per-episode task strings + for free via LeRobot's `task_index`, which its `modality.json` references. + +## Design + +### 1. Per-episode task instruction + +Add one new optional field to `ExportConfig` (`config.py`): + +```python +task_field: str | None = Field( + default=None, + description=( + "meta.json field name (e.g. 'task') holding this episode's " + "natural-language task instruction. When set, each episode's LeRobot " + "'task' is read from its own meta.json instead of the fixed " + "task_description. Falls back to task_description when the field is " + "missing for a given episode, or when no local meta.json is " + "available (e.g. exporting from catalog_url). Requires local " + "rrd_paths export, like episode_metadata." + ), +) +``` + +Default `None` — behavior is byte-for-byte identical to today (one constant +`task_description` for the whole dataset) unless a user opts in. + +**Loading**: `exporter.py`'s existing `_load_episode_metadata` already opens +each recording's sibling `meta.json` once per recording. Extend it to also +pull `task_field`'s value in that same read (avoid opening the file twice), +returning it alongside the metadata dict rather than adding a second loader. + +**Threading through**: `Episode` (in `episode_sampler.py`) gets one new +optional attribute, `task: str | None`, set in `exporter.py` next to where +`episode.extra_metadata` is already set today — `episode.task` is the +resolved per-episode string (meta.json value if present, else +`config.task_description`, with a warning on fallback — mirroring the +existing missing-field warning for `episode_metadata`). + +**Writing**: `heads/lerobot.py:_sample_to_frame` takes the resolved task +string as a parameter (from the enclosing `episode.task` in `write_episode`) +instead of always reading `self.config.task_description` directly. LeRobot +natively supports a per-episode-varying `task` column (that's what its +`tasks`/`task_index` table exists for) — no changes needed to +`initialize()`/`finalize()`. + +### 2. Generalize `episode_metadata` + +In `exporter.py`: + +- `_load_episode_metadata`'s return type becomes + `dict[str, dict[str, Any]]` (was `dict[str, dict[str, float]]`) — the + underlying read (`meta[f]`) already preserves whatever JSON type is + present; only the type annotation was wrong. +- `episode.extra_metadata = {f: found.get(f, 0.0) for f in ...}` → + `{f: found.get(f) for f in ...}` (default `None`, not `0.0`). Update the + existing "filled with 0.0" warning log to say "filled with null". + +No change needed in `heads/lerobot.py:_write_episode_metadata_columns` — +pandas/pyarrow already handle a `None`-containing or string-typed column +into Parquet without modification. + +## Data flow (updated) + +``` +meta.json (per recording, optional) + │ + ├─ task_field value ──────────► episode.task ──► frame["task"] (per sample) + │ (falls back to config.task_description) + │ + └─ episode_metadata values ───► episode.extra_metadata ──► meta/episodes/*.parquet + (any JSON scalar; missing → null) +``` + +Neither path does anything when its config field is unset/empty — a bare +imitation-learning dataset with no `meta.json` metadata at all exports +exactly as today. + +## Error handling + +- Missing `meta.json` entirely (no local rrd_paths, i.e. `catalog_url` + export): both `task_field` and `episode_metadata` are skipped with a + warning (existing behavior for `episode_metadata`, extended to + `task_field`). +- `task_field` missing from one recording's `meta.json` while present in + others: that episode falls back to `config.task_description`, with a + per-episode warning — never a hard failure (matches how a missing + `episode_metadata` field is already handled: warn and fill, not abort). + +## Compatibility with standard LeRobot / other datasets + +Verified against the installed `lerobot` package, not just assumed: + +- `frame["task"]` → `meta/tasks.parquet` + `task_index` is the mechanism + every standard LeRobot dataset uses for language conditioning, and what + `aggregate_datasets`'s schema check (`validate_all_metadata` / + `features_equal_for_merge`) and every stock policy config (ACT, diffusion, + pi0, SmolVLA) actually read. Using it (rather than inventing a parallel + column) is what makes this dataset minglable with datasets from other + sources for training, not just internally consistent. +- This `lerobot` version also has a separate, newer `language_persistent`/ + `language_events` schema (structured subtask/plan/memory/VQA annotations + with roles and timestamps) — but it's populated by a standalone offline + annotation tool (`lerobot_annotate` / `steerable_pipeline`), not by + exporters, and essentially no external dataset will have it. Out of scope + here; revisit only if a target policy specifically requires it. +- `episode_metadata`'s extra `meta/episodes/*.parquet` columns never + participate in `aggregate_datasets`'s validation (which only compares + frame-level `features`), so they can't affect mixing with other datasets + for training — they're purely additive/inert to any standard tooling. + +**Cross-checked against real published community datasets, not just the +local package:** + +- [DROID](https://github.com/google-deepmind/open_x_embodiment), a major + Open X-Embodiment dataset, stores its primary instruction as a per-episode + `language_instruction` string feeding the same `task`/`tasks.jsonl`/ + `task_index` mechanism this design uses — confirming it as the real-world + standard for the *portable* instruction. +- [AgiBot World 2026](https://huggingface.co/datasets/agibot-world/AgiBotWorld2026), + a large published LeRobot-format dataset, layers custom per-episode + annotations (subtask instructions, object bounding boxes) in + `meta/info.json` *alongside*, not replacing, the standard `tasks.jsonl` + mechanism — explicitly to stay compatible with standard training + pipelines. Same shape as this design: one portable `task` string plus + optional non-standard extra metadata that doesn't affect portability. + +## Testing + +- `tests/test_export.py`: unit tests for `_load_episode_metadata` covering + (a) a string field (`cube_color`) round-tripping without corruption, (b) a + missing field defaulting to `None` (not `0.0`), (c) `task_field` present → + used verbatim, (d) `task_field` absent for one episode → falls back to + `task_description` with a warning, (e) neither field configured → output + identical to current behavior (regression guard for the optionality + requirement). +- `tests/test_export.py` (LeRobotHead-level): an episode written with a + distinct `task_field` value produces that string in the frame's `task`, + and multiple episodes with different task strings both land in the + dataset's task table (proves LeRobot's per-episode task mechanism is being + used correctly, not just a per-dataset constant renamed). From 7b6238e4ec4f84f9c0f689138e2ee8509dd0b887 Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Tue, 11 Aug 2026 11:29:56 +0200 Subject: [PATCH 08/12] feat: per-episode task instructions + generalized episode metadata Adds two optional, backward-compatible exporter capabilities for VLA training datasets (LeRobot export head): - config.task_field: names a meta.json field (e.g. "task") holding each episode's own natural-language instruction. Resolved per episode and written via LeRobot's native frame["task"]/task_index mechanism, so different episodes can carry different instructions without inventing a parallel column. Falls back to the existing task_description when unset or missing for a given episode. - episode_metadata now accepts any JSON scalar (str/float/int/bool) instead of float-only, and defaults a missing field to null instead of the previously-wrong 0.0 default (e.g. a string field like cube_color). Both default to today's exact behavior when unset. Verified end-to-end against a real recording (task instruction and cube-position metadata round-tripped correctly into the exported dataset). Co-Authored-By: Claude Sonnet 5 --- docs/export-guide.md | 5 +- src/nova_export/export/config.py | 15 ++ src/nova_export/export/episode_sampler.py | 3 +- src/nova_export/export/exporter.py | 82 +++++++-- src/nova_export/export/heads/lerobot.py | 10 +- tests/test_export.py | 215 ++++++++++++++++++++++ 6 files changed, 310 insertions(+), 20 deletions(-) diff --git a/docs/export-guide.md b/docs/export-guide.md index cc1ab55..0b0961f 100644 --- a/docs/export-guide.md +++ b/docs/export-guide.md @@ -66,10 +66,13 @@ directly and `--recordings-dir` is ignored. | `cameras` | list[object] | `[]` | Camera streams → `observation.images.`. Each may set `width`/`height` to resize. See [Cameras](#cameras--resizing). | | `trimming` | object | `all_present` | How episode start/end bounds are chosen. See [Trimming](#trimming-the-important-part). | | `max_episode_duration_s` | float \| null | `null` (no limit) | Reject a segment if its _raw_ recording span exceeds this many seconds — unrelated to trimming. See [Rejecting stuck or left-running recordings](#rejecting-stuck-or-left-running-recordings). | -| `task_description` | string | `"task"` | Natural-language task label written to every frame. | +| `task_description` | string | `"task"` | Fallback task label written to every frame when `task_field` is unset (or its meta.json field is missing for a given episode). | +| `task_field` | string \| null | `null` | meta.json field name (e.g. `"task"`) holding each episode's own natural-language instruction — lets `task` vary per episode instead of being fixed dataset-wide. Falls back to `task_description`. Requires local export (same as `episode_metadata`). | | `dataset_id` | string | `nova/dataset` | Dataset identifier — the LeRobot `repo_id` (also used for viz and Hugging Face push). | | `version` | int | `1` | Config schema version. Leave at `1`. | +`episode_metadata` values may be any JSON scalar type (string, number, boolean) — not float-only. A field missing from a given episode's `meta.json` is filled with `null` in that episode's row. + ## Formats - **`lerobot_v3`** — a LeRobot v3.0 dataset (Parquet + MP4 + metadata). Use this diff --git a/src/nova_export/export/config.py b/src/nova_export/export/config.py index aa0c78f..d4e6260 100644 --- a/src/nova_export/export/config.py +++ b/src/nova_export/export/config.py @@ -158,6 +158,21 @@ class ExportConfig(BaseModel): description="Task label written to the dataset", ) + task_field: str | None = Field( + default=None, + description=( + "meta.json field name (e.g. 'task') holding this episode's " + "natural-language task instruction. When set, each episode's " + "LeRobot 'task' is read from its own meta.json instead of the " + "fixed task_description, mirroring how LeRobot's task/task_index " + "mechanism is meant to vary per episode. Falls back to " + "task_description when the field is missing for a given episode, " + "or when no local meta.json is available (e.g. exporting from " + "catalog_url). Requires local rrd_paths export, like " + "episode_metadata." + ), + ) + dataset_id: str = Field( default="nova/dataset", description="Dataset identifier (repo_id for LeRobot, dataset name for Groot)", diff --git a/src/nova_export/export/episode_sampler.py b/src/nova_export/export/episode_sampler.py index 71eacf7..bc6a85e 100644 --- a/src/nova_export/export/episode_sampler.py +++ b/src/nova_export/export/episode_sampler.py @@ -63,7 +63,8 @@ class Episode: segment_id: str episode_index: int samples: list[Sample] - extra_metadata: dict[str, float] | None = None + extra_metadata: dict[str, Any] | None = None + task: str | None = None @property def num_frames(self) -> int: diff --git a/src/nova_export/export/exporter.py b/src/nova_export/export/exporter.py index ec008ea..854eecd 100644 --- a/src/nova_export/export/exporter.py +++ b/src/nova_export/export/exporter.py @@ -23,6 +23,7 @@ import json from collections.abc import Callable, Generator from pathlib import Path +from typing import Any import rerun as rr from loguru import logger @@ -200,34 +201,72 @@ def _validate_sources(dataset, config: ExportConfig, segment_id: str) -> None: ) +def _resolve_extra_metadata(found: dict[str, Any], fields: list[str]) -> dict[str, Any]: + """Build one episode's extra_metadata dict from its meta.json values. + + A field missing from `found` becomes None, not 0.0 — 0.0 was only + correct by accident for numeric fields and actively wrong for a field + like cube_color ("purple"). + """ + return {f: found.get(f) for f in fields} + + +def _resolve_task( + found: dict[str, Any], task_field: str | None, task_description: str +) -> str: + """Resolve one episode's task string. + + task_field's value from meta.json when set and present; task_description + otherwise (unset task_field, or the field missing from this episode's + meta.json) — the same fallback either way, so a per-recording gap in + metadata degrades to today's dataset-wide constant rather than failing. + """ + if not task_field: + return task_description + value = found.get(task_field) + return task_description if value is None else str(value) + + def _load_episode_metadata( - rrd_paths: list[Path] | None, fields: list[str] -) -> dict[str, dict[str, float]]: - """Read episode_metadata fields from each recording's sibling meta.json. + rrd_paths: list[Path] | None, + fields: list[str], + task_field: str | None = None, +) -> dict[str, dict[str, Any]]: + """Read episode_metadata fields (and optionally task_field) from each + recording's sibling meta.json. Keyed by segment_id, which for local rrd_paths exports is exactly the recording's directory name (//recording.rrd) — the same recording_id the collector assigns and rerun uses as the - segment ID, so no separate ID plumbing is needed. + segment ID, so no separate ID plumbing is needed. Values are whatever + JSON scalar type meta.json holds (str/float/int/bool) — not float-only. + + When task_field is set, its value (if present) is included in the + per-segment dict under its own key, alongside the requested + episode_metadata fields — one meta.json read serves both. """ - if not fields: + if not fields and not task_field: return {} if not rrd_paths: logger.warning( - "episode_metadata {} configured but exporting from catalog_url " - "(no local meta.json available) — skipping", + "episode_metadata {} / task_field {!r} configured but exporting " + "from catalog_url (no local meta.json available) — skipping", fields, + task_field, ) return {} - result: dict[str, dict[str, float]] = {} + result: dict[str, dict[str, Any]] = {} for rrd_path in rrd_paths: meta_path = rrd_path.parent / "meta.json" if not meta_path.is_file(): continue meta = json.loads(meta_path.read_text()) segment_id = rrd_path.parent.name - result[segment_id] = {f: meta[f] for f in fields if f in meta} + entry: dict[str, Any] = {f: meta[f] for f in fields if f in meta} + if task_field and task_field in meta: + entry[task_field] = meta[task_field] + result[segment_id] = entry return result @@ -393,7 +432,7 @@ def export_recordings( head = _create_export_head(config, output_dir) episode_metadata_by_segment = _load_episode_metadata( - rrd_paths, config.episode_metadata + rrd_paths, config.episode_metadata, config.task_field ) # For the max_episode_duration_s safety check: fetch the dataset's @@ -488,19 +527,32 @@ def export_recordings( ) continue + found = episode_metadata_by_segment.get(segment_id, {}) + if config.episode_metadata: - found = episode_metadata_by_segment.get(segment_id, {}) missing = [f for f in config.episode_metadata if f not in found] if missing: logger.warning( - "Episode {} ({}): meta.json missing {} — filled with 0.0", + "Episode {} ({}): meta.json missing {} — filled with null", episode_id, segment_id[:8], missing, ) - episode.extra_metadata = { - f: found.get(f, 0.0) for f in config.episode_metadata - } + episode.extra_metadata = _resolve_extra_metadata( + found, config.episode_metadata + ) + + if config.task_field and found.get(config.task_field) is None: + logger.warning( + "Episode {} ({}): meta.json missing task_field {!r} " + "— falling back to task_description", + episode_id, + segment_id[:8], + config.task_field, + ) + episode.task = _resolve_task( + found, config.task_field, config.task_description + ) # Check if episode has samples if not episode.samples: diff --git a/src/nova_export/export/heads/lerobot.py b/src/nova_export/export/heads/lerobot.py index c242946..6a03daa 100644 --- a/src/nova_export/export/heads/lerobot.py +++ b/src/nova_export/export/heads/lerobot.py @@ -157,9 +157,11 @@ def write_episode(self, episode: Episode) -> bool: episode.extra_metadata ) + task = episode.task if episode.task is not None else self.config.task_description + try: for sample in tqdm(episode.samples, desc="Frames", leave=False): - frame = self._sample_to_frame(sample) + frame = self._sample_to_frame(sample, task) self._dataset.add_frame(frame) self._dataset.save_episode() @@ -227,11 +229,13 @@ def _write_episode_metadata_columns(self) -> None: len(episodes_files), ) - def _sample_to_frame(self, sample: Sample) -> dict[str, Any]: + def _sample_to_frame(self, sample: Sample, task: str) -> dict[str, Any]: """Convert a Sample to a LeRobot frame dict. Args: sample: Sample to convert. + task: Resolved task string for this sample's episode (either + Episode.task, when set, or config.task_description). Returns: Frame dict for LeRobotDataset.add_frame(). @@ -246,7 +250,7 @@ def _sample_to_frame(self, sample: Sample) -> dict[str, Any]: frame["observation.state"] = sample.state # Task - frame["task"] = self.config.task_description + frame["task"] = task # Images for cam_name, img_array in sample.images.items(): diff --git a/tests/test_export.py b/tests/test_export.py index dd50dba..47715c0 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -171,6 +171,23 @@ def create_test_episode( ) +# ============================================================================= +# ExportConfig Tests +# ============================================================================= + + +class TestExportConfigTaskField: + """Tests for ExportConfig.task_field.""" + + def test_task_field_defaults_to_none(self): + config = ExportConfig(fps=15) + assert config.task_field is None + + def test_task_field_can_be_set(self): + config = ExportConfig(fps=15, task_field="task") + assert config.task_field == "task" + + # ============================================================================= # FrameCache Tests # ============================================================================= @@ -709,6 +726,75 @@ def test_write_episode(self, mock_dataset_cls): assert mock_dataset.add_frame.call_count == 5 mock_dataset.save_episode.assert_called_once() + @patch("lerobot.datasets.lerobot_dataset.LeRobotDataset") + def test_write_episode_uses_per_episode_task(self, mock_dataset_cls): + """When Episode.task is set, every frame's 'task' must use it — + not the dataset-wide config.task_description.""" + mock_dataset = MagicMock() + mock_dataset_cls.create.return_value = mock_dataset + + config = ExportConfig(fps=15, task_description="fallback_task") + + with tempfile.TemporaryDirectory() as tmpdir: + head = LeRobotHead(config, Path(tmpdir) / "output") + head.initialize({"action": {"dtype": "float32", "shape": (7,)}}) + + episode = create_test_episode(num_samples=3) + episode.task = "Pick the purple cube up." + head.write_episode(episode) + + for call in mock_dataset.add_frame.call_args_list: + frame = call.args[0] + assert frame["task"] == "Pick the purple cube up." + + @patch("lerobot.datasets.lerobot_dataset.LeRobotDataset") + def test_write_episode_falls_back_to_task_description(self, mock_dataset_cls): + """When Episode.task is unset (None), fall back to config.task_description — + this is the byte-for-byte-identical-to-today path.""" + mock_dataset = MagicMock() + mock_dataset_cls.create.return_value = mock_dataset + + config = ExportConfig(fps=15, task_description="fallback_task") + + with tempfile.TemporaryDirectory() as tmpdir: + head = LeRobotHead(config, Path(tmpdir) / "output") + head.initialize({"action": {"dtype": "float32", "shape": (7,)}}) + + episode = create_test_episode(num_samples=3) # episode.task defaults to None + + head.write_episode(episode) + + for call in mock_dataset.add_frame.call_args_list: + frame = call.args[0] + assert frame["task"] == "fallback_task" + + @patch("lerobot.datasets.lerobot_dataset.LeRobotDataset") + def test_different_episodes_can_have_different_tasks(self, mock_dataset_cls): + """Two episodes with distinct task strings both write their own value — + proves this is LeRobot's per-episode task mechanism, not a renamed + per-dataset constant.""" + mock_dataset = MagicMock() + mock_dataset_cls.create.return_value = mock_dataset + + config = ExportConfig(fps=15) + + with tempfile.TemporaryDirectory() as tmpdir: + head = LeRobotHead(config, Path(tmpdir) / "output") + head.initialize({"action": {"dtype": "float32", "shape": (7,)}}) + + episode_a = create_test_episode(segment_id="a", num_samples=1) + episode_a.task = "Task A" + episode_b = create_test_episode(segment_id="b", num_samples=1) + episode_b.task = "Task B" + + head.write_episode(episode_a) + head.write_episode(episode_b) + + frame_a = mock_dataset.add_frame.call_args_list[0].args[0] + frame_b = mock_dataset.add_frame.call_args_list[1].args[0] + assert frame_a["task"] == "Task A" + assert frame_b["task"] == "Task B" + @patch("lerobot.datasets.lerobot_dataset.LeRobotDataset") def test_write_empty_episode(self, mock_dataset_cls): """Test writing empty episode returns False.""" @@ -886,6 +972,135 @@ def test_duplicate_source_queried_once_and_shared(self): ) +class TestResolveExtraMetadata: + """Tests for exporter._resolve_extra_metadata.""" + + def test_present_fields_pass_through_any_type(self): + from nova_export.export.exporter import _resolve_extra_metadata + + found = {"cube_x_mm": -324.15, "cube_color": "purple"} + + result = _resolve_extra_metadata(found, ["cube_x_mm", "cube_color"]) + + assert result == {"cube_x_mm": -324.15, "cube_color": "purple"} + + def test_missing_field_defaults_to_none_not_zero(self): + from nova_export.export.exporter import _resolve_extra_metadata + + found = {"cube_x_mm": 1.5} # cube_color absent + + result = _resolve_extra_metadata(found, ["cube_x_mm", "cube_color"]) + + assert result == {"cube_x_mm": 1.5, "cube_color": None} + + +class TestResolveTask: + """Tests for exporter._resolve_task.""" + + def test_task_field_unset_uses_task_description(self): + from nova_export.export.exporter import _resolve_task + + result = _resolve_task({}, None, "fallback_task") + + assert result == "fallback_task" + + def test_task_field_present_used_verbatim(self): + from nova_export.export.exporter import _resolve_task + + found = {"task": "Pick the purple cube up."} + + result = _resolve_task(found, "task", "fallback_task") + + assert result == "Pick the purple cube up." + + def test_task_field_missing_from_found_falls_back(self): + from nova_export.export.exporter import _resolve_task + + found = {"other_field": 1} # no "task" key + + result = _resolve_task(found, "task", "fallback_task") + + assert result == "fallback_task" + + def test_task_field_present_but_null_falls_back(self): + # meta.json had `"task": null` — the key exists but its value is + # None. This must fall back exactly like a missing key, and the + # per-episode warning in export_recordings uses + # `found.get(config.task_field) is None` (not `not in found`) so it + # fires on this case too. + from nova_export.export.exporter import _resolve_task + + found = {"task": None} + + result = _resolve_task(found, "task", "fallback_task") + + assert result == "fallback_task" + assert found.get("task") is None # same predicate export_recordings warns on + + +class TestLoadEpisodeMetadata: + """Tests for exporter._load_episode_metadata.""" + + def test_string_field_round_trips(self, tmp_path): + from nova_export.export.exporter import _load_episode_metadata + + rec_dir = tmp_path / "04cb4f25d3ef" + rec_dir.mkdir() + (rec_dir / "meta.json").write_text( + '{"cube_color": "purple", "cube_x_mm": -324.15}' + ) + rrd_path = rec_dir / "recording.rrd" + rrd_path.touch() + + result = _load_episode_metadata([rrd_path], ["cube_color", "cube_x_mm"]) + + assert result["04cb4f25d3ef"]["cube_color"] == "purple" + assert result["04cb4f25d3ef"]["cube_x_mm"] == -324.15 + + def test_missing_field_simply_absent_from_result(self, tmp_path): + from nova_export.export.exporter import _load_episode_metadata + + rec_dir = tmp_path / "rec01" + rec_dir.mkdir() + (rec_dir / "meta.json").write_text('{"cube_x_mm": 1.5}') + rrd_path = rec_dir / "recording.rrd" + rrd_path.touch() + + result = _load_episode_metadata([rrd_path], ["cube_x_mm", "cube_color"]) + + assert result["rec01"] == {"cube_x_mm": 1.5} + assert "cube_color" not in result["rec01"] + + def test_task_field_value_included_in_result(self, tmp_path): + from nova_export.export.exporter import _load_episode_metadata + + rec_dir = tmp_path / "04cb4f25d3ef" + rec_dir.mkdir() + (rec_dir / "meta.json").write_text( + '{"task": "Pick the purple cube up.", "cube_x_mm": 1.0}' + ) + rrd_path = rec_dir / "recording.rrd" + rrd_path.touch() + + result = _load_episode_metadata([rrd_path], ["cube_x_mm"], task_field="task") + + assert result["04cb4f25d3ef"]["task"] == "Pick the purple cube up." + assert result["04cb4f25d3ef"]["cube_x_mm"] == 1.0 + + def test_task_field_none_does_not_add_task_key(self, tmp_path): + from nova_export.export.exporter import _load_episode_metadata + + rec_dir = tmp_path / "rec01" + rec_dir.mkdir() + (rec_dir / "meta.json").write_text('{"task": "unused", "cube_x_mm": 1.0}') + rrd_path = rec_dir / "recording.rrd" + rrd_path.touch() + + result = _load_episode_metadata([rrd_path], ["cube_x_mm"]) # task_field omitted + + assert result["rec01"] == {"cube_x_mm": 1.0} + + # ============================================================================= # Source validation (helpful error when a config source is missing) # ============================================================================= From fcaa250b6aca9a2769e43e4814a90e56a29b31f9 Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Tue, 11 Aug 2026 12:25:32 +0200 Subject: [PATCH 09/12] chore: point pipeline.sh defaults at pick_and_place_sim_20260810_201248 Reuses the existing lerobot_export.json in pick_and_place_imitation_learning (no new config file) and this machine's paths, so the pipeline can be run with no PIPELINE_* env vars for this dataset. Co-Authored-By: Claude Sonnet 5 --- tools/pipeline.sh | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tools/pipeline.sh b/tools/pipeline.sh index bad3a5b..f1a5fd5 100755 --- a/tools/pipeline.sh +++ b/tools/pipeline.sh @@ -13,14 +13,14 @@ if [[ -n "${PIPELINE_REMOTE_DIRS:-}" ]]; then IFS=':' read -r -a REMOTE_DIRS <<< "$PIPELINE_REMOTE_DIRS" else REMOTE_DIRS=( - "/mnt/data/sebastian/raw_datasets/dryrun_pose_algo_check" + "/mnt/data/sebastian/raw_datasets/pick_and_place_sim_20260810_201248" ) fi -WATCH_DIR="${PIPELINE_WATCH_DIR:-/mnt/data/sebastian/raw_datasets/dryrun_pose_algo_check}" +WATCH_DIR="${PIPELINE_WATCH_DIR:-/home/sebi/ws/Data/raw_data/pick_and_place_sim_20260810_201248}" -NOVA_CLI_DIR="${PIPELINE_NOVA_CLI_DIR:-/home/intern/ws/nova-data-cli}" -EXPORT_CONFIG="${PIPELINE_EXPORT_CONFIG:-/home/intern/ws/pick_and_place_imitation_learning/data_collection/configs/lerobot_export.json}" -EXPORT_ROOT="${PIPELINE_EXPORT_ROOT:-/mnt/data/sebastian/lerobot_datasets/dryrun_pose_algo_check}" +NOVA_CLI_DIR="${PIPELINE_NOVA_CLI_DIR:-/home/sebi/ws/nova-data-cli}" +EXPORT_CONFIG="${PIPELINE_EXPORT_CONFIG:-/home/sebi/ws/pick_and_place_imitation_learning/data_collection/configs/lerobot_export.json}" +EXPORT_ROOT="${PIPELINE_EXPORT_ROOT:-/home/sebi/ws/Data/lerobot_datasets/pick_and_place_sim_20260810_201248}" read -r -a EXPORT_CLI_CMD <<< "${PIPELINE_EXPORT_CLI:-uv run nova-data-cli}" # swap in a stub for tests CHUNK="${PIPELINE_CHUNK:-8}" From 4dccc65ebc3a8f44d8bf96a4f47344f270db2de2 Mon Sep 17 00:00:00 2001 From: Sebastian Dominguez Date: Tue, 11 Aug 2026 12:25:42 +0200 Subject: [PATCH 10/12] fix: make Ctrl+C actually stop pipeline.sh setsid detaches the supervisor from the controlling terminal, so Ctrl+C's SIGINT never reached it -- the shutdown trap was correctly written but never invoked. An outer wrapper now stays attached to the terminal and forwards INT/TERM into the detached process group. Also fixes a second bug this surfaced: the shutdown trap's own `kill -- -$$` re-signals itself (same pgid), re-entering the same trap forever instead of exiting. Trap now disarms itself before killing. Co-Authored-By: Claude Sonnet 5 --- tools/pipeline.sh | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/tools/pipeline.sh b/tools/pipeline.sh index f1a5fd5..f9adec6 100755 --- a/tools/pipeline.sh +++ b/tools/pipeline.sh @@ -337,10 +337,26 @@ role_worker() { # ---- supervisor ----------------------------------------------------------- role_supervisor() { - # setsid makes this the leader of a fresh process group so `kill -- -$$` on - # shutdown reaches every descendant (see AGENT.md). + # setsid makes the real supervisor (below) the leader of a fresh + # session/process group so `kill -- -$$` on shutdown reliably reaches + # every descendant (workers, nova-data-cli, ffmpeg) regardless of how this + # script was invoked -- e.g. piped, where the default process group would + # otherwise be shared with unrelated commands (see AGENT.md). + # + # setsid also detaches from the controlling terminal, which has a side + # effect: Ctrl+C's SIGINT is delivered by the terminal driver to whatever + # process group the terminal currently has registered as foreground, and + # a setsid'd process is never in that group -- so Ctrl+C would otherwise + # go nowhere, even though the trap below is correctly set up to handle it. + # This outer wrapper stays attached to the terminal specifically to catch + # that Ctrl+C/SIGINT and forward it into the detached process group, so an + # interactive `bash tools/pipeline.sh` still responds to Ctrl+C normally. if [[ -z "${PIPELINE_RESPAWNED:-}" ]]; then - exec env PIPELINE_RESPAWNED=1 setsid "$0" supervisor --mode "$MODE" + PIPELINE_RESPAWNED=1 setsid "$0" supervisor --mode "$MODE" & + local inner_pid=$! + trap 'kill -TERM -- "-$inner_pid" 2>/dev/null' INT TERM + wait "$inner_pid" + exit $? fi exec 9>"$LOCK" @@ -366,7 +382,9 @@ role_supervisor() { log "sizing: $(nproc) cores, $(awk '/MemTotal/{printf "%.1fGB", $2/1024/1024}' /proc/meminfo) RAM, MEM_TARGET_FRACTION=$MEM_TARGET_FRACTION, WORKER_MEM_ESTIMATE_MB=$WORKER_MEM_ESTIMATE_MB -> starting at $(compute_workers) workers, will scale with available memory" - trap 'log "shutting down"; kill -- -$$ 2>/dev/null || true' INT TERM + # kill -- -$$ also signals this process itself (same pgid) -> re-enters + # this same trap before reaching exit, looping forever. Disarm first. + trap 'trap - INT TERM; log "shutting down"; kill -- -$$ 2>/dev/null || true; exit 130' INT TERM # Tee each role's output to both terminal (prefixed) and its log file. ( "$0" acquire --mode "$MODE" 2>&1 | sed -u 's/^/[acquire] /' | tee -a "$LOGS/acquire.log" ) & From 5fc15af2e4e4a19e943f3b765a009910e2f9b413 Mon Sep 17 00:00:00 2001 From: Shamreen Tabassum Date: Sat, 15 Aug 2026 13:16:26 +0200 Subject: [PATCH 11/12] fix: eliminate export-time episode truncation and camera-ahead-of-action violations Two real defects found and fixed in the export pipeline, verified against real recordings from pick_and_place_v4_1000: 1. signal_change trimming (_apply_trimming) compared its threshold against a per-sample joint-position delta, but the default threshold (0.01) sat above the 99th percentile of real per-sample motion for this robot/log rate. As a result, trimming matched almost nothing but the fastest samples (typically only the final retreat), collapsing episodes to a tiny fragment (as little as 1.6s of a true ~40s Leg B). Fixed the example configs to default to all_present trimming, lowered the threshold as a fallback default, and added a guard that warns loudly (with the observed signal statistics) whenever signal_change trimming keeps less than half the raw episode span, instead of silently shipping a fragment. 2. _query_action_state resampled action/state onto the export's fixed-rate grid via nearest-neighbor, which could select a sample from AFTER the grid timestamp. Because /actions_target has real dropouts (tens to hundreds of ms), this let some exported frames pair a camera timestamp with a numerically later action sample - i.e. camera-ahead-of-action in the exported data, undermining the action-lead invariant the collector otherwise guarantees at recording time. Replaced nearest-neighbor with strictly causal selection (latest sample at or before grid time), which cannot select a future sample. Verified on 5 real re-exported episodes: violations dropped from 425/9168 frame-camera pairs (4.6%, up to 331.7ms) to 3/9168 (0.03%, all attributable to genuine raw camera dropouts, not the export logic). Re-exporting 4-5 real episodes before/after confirms both fixes: exported duration now matches the raw HOME-to-HOME Leg B span (previously as low as 3.9% of it), and the exported action column exactly matches the raw /actions_target value at its causal grid selection. Test suite: 102 passed, 3 skipped. --- docs/export-guide.md | 15 +++++- examples/groot_export.json | 14 +++--- examples/lerobot_export.json | 16 +++--- examples/lerobot_export_resized.json | 14 +++--- src/nova_export/export/episode_sampler.py | 59 ++++++++++++++++------- 5 files changed, 81 insertions(+), 37 deletions(-) diff --git a/docs/export-guide.md b/docs/export-guide.md index 0b0961f..1a4d1ab 100644 --- a/docs/export-guide.md +++ b/docs/export-guide.md @@ -197,6 +197,17 @@ depends on the signal's units and noise floor: Start at `0.01` and raise it only if idle time is leaking in; if episodes come out suspiciously short, your threshold is above the real motion and should come down. +**`threshold` is per *consecutive sample*, not total displacement**, so the right +value depends on the source's sample rate as much as on its units. A slow arm +logged at 64 ms may never move more than ~0.01 rad between two samples, in which +case `threshold: 0.01` trims almost the entire episode away. The "no change +exceeds it" fallback above does *not* save you here — a handful of samples still +cross, so the episode collapses to a second or two instead of falling back. The +export logs a warning whenever trimming keeps less than half the raw span; treat +it as a signal to lower `threshold` or switch to `all_present`. When the action +stream itself only exists while the task is being commanded, `all_present` +already trims the idle lead-in for free and is the safer choice. + ### Modes compared ![All trim modes compared](img/all_modes_compared.png) @@ -228,7 +239,9 @@ check, so this doesn't need to be tight. - **Just want everything recorded?** `all_present` (default). - **A signal cleanly brackets the task?** `signal_presence` on that source. - **Need to cut idle lead-in/out automatically?** `signal_change` on a motion - signal (e.g. `joint_positions`), `threshold` ≈ `0.01`, `tail_ms` ≈ `500`. + signal (e.g. `joint_positions`), `threshold` ≈ `0.01`, `tail_ms` ≈ `500` — but + check the threshold against your source's actual inter-sample motion first + (see [Tuning `threshold`](#signal_change)). - **Dataset too big / training input smaller?** Set camera `width`/`height`. - **A few episodes are way longer than the rest (stuck sensor, forgotten recording)?** Set `max_episode_duration_s` to drop them. diff --git a/examples/groot_export.json b/examples/groot_export.json index 5ddd808..608f0de 100644 --- a/examples/groot_export.json +++ b/examples/groot_export.json @@ -3,8 +3,13 @@ "format": "groot", "fps": 15, "index_column": "canonical_time", - "action": ["actions_target"], - "state": ["joint_positions", "gripper"], + "action": [ + "actions_target" + ], + "state": [ + "joint_positions", + "gripper" + ], "cameras": [ { "source": "cam_flange" @@ -17,10 +22,7 @@ } ], "trimming": { - "mode": "signal_change", - "source": "joint_positions", - "threshold": 0.01, - "tail_ms": 500 + "mode": "all_present" }, "task_description": "pick-and-place", "dataset_id": "nova/pick-and-place" diff --git a/examples/lerobot_export.json b/examples/lerobot_export.json index 067e98a..69eadf7 100644 --- a/examples/lerobot_export.json +++ b/examples/lerobot_export.json @@ -3,8 +3,13 @@ "format": "lerobot_v3", "fps": 15, "index_column": "canonical_time", - "action": ["actions_target"], - "state": ["joint_positions", "gripper"], + "action": [ + "actions_target" + ], + "state": [ + "joint_positions", + "gripper" + ], "cameras": [ { "source": "cam_flange" @@ -17,11 +22,8 @@ } ], "trimming": { - "mode": "signal_change", - "source": "joint_positions", - "threshold": 0.01, - "tail_ms": 500 + "mode": "all_present" }, "task_description": "pick-and-place", "dataset_id": "nova/pick-and-place" -} \ No newline at end of file +} diff --git a/examples/lerobot_export_resized.json b/examples/lerobot_export_resized.json index 26f3dc4..f3036e2 100644 --- a/examples/lerobot_export_resized.json +++ b/examples/lerobot_export_resized.json @@ -3,8 +3,13 @@ "format": "lerobot_v3", "fps": 15, "index_column": "canonical_time", - "action": ["actions_target"], - "state": ["joint_positions", "gripper"], + "action": [ + "actions_target" + ], + "state": [ + "joint_positions", + "gripper" + ], "cameras": [ { "source": "cam_flange", @@ -23,10 +28,7 @@ } ], "trimming": { - "mode": "signal_change", - "source": "joint_positions", - "threshold": 0.01, - "tail_ms": 500 + "mode": "all_present" }, "task_description": "pick-and-place", "dataset_id": "nova/pick-and-place-resized" diff --git a/src/nova_export/export/episode_sampler.py b/src/nova_export/export/episode_sampler.py index bc6a85e..c717de9 100644 --- a/src/nova_export/export/episode_sampler.py +++ b/src/nova_export/export/episode_sampler.py @@ -30,6 +30,10 @@ if TYPE_CHECKING: from nova_export.export.config import ExportConfig +# signal_change trimming that keeps less than this share of the raw span is +# almost certainly a mis-tuned threshold rather than a genuinely short episode. +_TRIM_SUSPICIOUS_FRACTION = 0.5 + @dataclass class Sample: @@ -416,6 +420,32 @@ def _apply_trimming( ) # index into times (diff shifts by 1) trim_end = int(times[last_active_idx]) + cfg.tail_ms * 1_000_000 + # A threshold set above the signal's real inter-sample motion does + # not always trip the "no activity at all" fallback above: a few + # noise/fast-move samples can still cross it and silently collapse a + # 40s episode to a second or two. Warn when the kept span is a small + # fraction of the raw one — that is the signature of too high a + # threshold, not of a genuinely short episode. + raw_span = raw_end_ns - raw_start_ns + kept_span = trim_end - trim_start + if raw_span > 0 and kept_span < _TRIM_SUSPICIOUS_FRACTION * raw_span: + logger.warning( + "Trimming kept only {:.1f}s of {:.1f}s in segment {} " + "({} of {} '{}' samples exceeded threshold={}). The threshold " + "is likely above this signal's real inter-sample motion " + "(median change {:.4g}, max {:.4g}) — lower it, or use " + "mode='all_present'.", + kept_span / 1e9, + raw_span / 1e9, + segment_id[:8], + len(active_indices), + len(diffs), + cfg.source, + cfg.threshold, + float(np.median(diffs)), + float(np.max(diffs)), + ) + return (trim_start, trim_end) return (raw_start_ns, raw_end_ns) @@ -523,27 +553,22 @@ def _query_action_state( ) raw_data[int(ts)] = {"action": action, "state": state} - # Resample to time_grid using nearest-neighbor lookup + # Resample to time_grid with a latest-at (causal) lookup: each grid + # point takes the most recent action/state at or before it. Nearest- + # neighbor would happily pick a *future* sample — with this action + # stream's dropouts (dt up to several hundred ms) that leaks commands + # issued after the observation they're paired with, i.e. exports an + # action that is "behind" its camera frame. Latest-at can never do that. sorted_query_ts = np.array(sorted(raw_data.keys()), dtype=np.int64) data: dict[int, dict[str, npt.NDArray[np.float32]]] = {} for target_ts in time_grid: - # Find nearest query timestamp - idx = np.searchsorted(sorted_query_ts, target_ts) - if idx == 0: - nearest_ts = sorted_query_ts[0] - elif idx >= len(sorted_query_ts): - nearest_ts = sorted_query_ts[-1] - else: - # Pick closest - if (target_ts - sorted_query_ts[idx - 1]) <= ( - sorted_query_ts[idx] - target_ts - ): - nearest_ts = sorted_query_ts[idx - 1] - else: - nearest_ts = sorted_query_ts[idx] - - data[int(target_ts)] = raw_data[int(nearest_ts)] + # searchsorted(side="right") - 1 is the last sample at or before + # target_ts; clamp at 0 for grid points preceding the first sample. + idx = int(np.searchsorted(sorted_query_ts, target_ts, side="right")) - 1 + latest_ts = sorted_query_ts[max(idx, 0)] + + data[int(target_ts)] = raw_data[int(latest_ts)] return data From 167b2538902114b70a70b24a4b2a6c75a1a19da5 Mon Sep 17 00:00:00 2001 From: Shamreen Tabassum Date: Wed, 19 Aug 2026 11:13:50 +0200 Subject: [PATCH 12/12] fix: never select a video frame earlier than its target grid timestamp get_frame_at/get_frame_index picked the nearest decoded frame to each grid timestamp, ties going to the earlier frame. That nearest-frame rule could still select a camera frame *before* the target time, re-opening the same camera-ahead/behind-of-action class of bug that 5fc15af just fixed on the action/state side. Both FrameCache.get_frame_index and VideoDecoder.decode_at now select the earliest frame at or after each target timestamp instead of the nearest one. Targets past the last decoded frame still clamp to it (now logged as a warning, since it means real target coverage was short). Co-Authored-By: Claude Sonnet 5 --- src/nova_export/export/episode_sampler.py | 9 +- src/nova_export/export/video_decoder.py | 100 +++++++++++----------- tests/test_export.py | 34 ++++---- 3 files changed, 76 insertions(+), 67 deletions(-) diff --git a/src/nova_export/export/episode_sampler.py b/src/nova_export/export/episode_sampler.py index c717de9..49ee8b6 100644 --- a/src/nova_export/export/episode_sampler.py +++ b/src/nova_export/export/episode_sampler.py @@ -2,7 +2,8 @@ This layer handles: - Building a fixed-rate time grid at target FPS -- Sampling video frames at grid timestamps (decoded streaming, grid-aligned) +- Sampling video frames at grid timestamps (decoded streaming, grid-aligned: + the earliest frame at or after each grid point, never an earlier one) - Querying action/state data at grid timestamps - Combining into unified Sample objects @@ -104,9 +105,9 @@ class EpisodeSampler: 2. Find the valid time range where all streams have data 3. Build a time grid at target FPS 4. Query action/state via fill_latest_at - 5. Decode each video stream once, keeping only the frames nearest to the - grid timestamps (already resized) — full-resolution frames are never - accumulated in memory + 5. Decode each video stream once, keeping only the earliest frame at or + after each grid timestamp (already resized) — full-resolution frames are + never accumulated in memory 6. Combine into Sample objects """ diff --git a/src/nova_export/export/video_decoder.py b/src/nova_export/export/video_decoder.py index eeba291..2aa7e08 100644 --- a/src/nova_export/export/video_decoder.py +++ b/src/nova_export/export/video_decoder.py @@ -5,7 +5,7 @@ - `load_packets()` + `decode_at()`: the memory-efficient path used by the export pipeline. Packets (compressed) are loaded first so time bounds are known before decoding, then the stream is decoded once and only the frames - nearest to the requested sample timestamps are kept (already resized). + at or after the requested sample timestamps are kept (already resized). - `decode_segment()`: decodes *all* frames into a FrameCache. Simple, but holds every decoded frame in memory — only suitable for short segments/tests. """ @@ -49,7 +49,8 @@ def resize_rgb( class FrameCache: """In-memory cache of decoded video frames with timestamp indexing. - Provides O(1) lookup of the nearest frame to any target timestamp. + Provides O(1) lookup of the earliest frame at or after any target + timestamp. """ frames: list[npt.NDArray[np.uint8]] = field(default_factory=list) @@ -77,34 +78,30 @@ def duration_s(self) -> float: return (self.end_ns - self.start_ns) / 1e9 def get_frame_at(self, target_ns: int) -> npt.NDArray[np.uint8] | None: - """Get the frame nearest to the target timestamp. + """Get the earliest frame at or after the target timestamp. Args: target_ns: Target timestamp in nanoseconds. Returns: - The nearest frame as HWC uint8 RGB array, or None if cache is empty. + The selected frame as HWC uint8 RGB array, or None if cache is + empty. """ idx = self.get_frame_index(target_ns) return self.frames[idx] if idx >= 0 else None def get_frame_index(self, target_ns: int) -> int: - """Get the index of the frame nearest to the target timestamp.""" + """Index of the earliest frame at or after the target timestamp. + + Never returns a frame *earlier* than the target — an observation can't + predate the action paired with it. Targets past the end of the stream + clamp to the last frame; -1 when the cache is empty. + """ if len(self.timestamps_ns) == 0: return -1 - idx = np.searchsorted(self.timestamps_ns, target_ns) - - if idx == 0: - return 0 - if idx >= len(self.timestamps_ns): - return len(self.timestamps_ns) - 1 - - if (target_ns - self.timestamps_ns[idx - 1]) <= ( - self.timestamps_ns[idx] - target_ns - ): - return idx - 1 - return idx + idx = int(np.searchsorted(self.timestamps_ns, target_ns, side="left")) + return min(idx, len(self.timestamps_ns) - 1) @dataclass @@ -408,15 +405,17 @@ def decode_at( target_timestamps_ns: npt.NDArray[np.int64], target_size: tuple[int, int] | None = None, ) -> FrameCache: - """Decode a segment, keeping only the frames nearest each target timestamp. + """Decode a segment, keeping the earliest frame at/after each target. The stream is decoded sequentially exactly once, but instead of caching - every decoded frame, each target timestamp is resolved to its nearest - frame on the fly (ties to the earlier frame, clamped at both ends — - the same selection rule as ``FrameCache.get_frame_at``). Only selected - frames are retained, already resized to ``target_size``, so peak memory - is one raw decoded frame plus the selected output frames. Decoding - stops early once every target is resolved. + every decoded frame, each target timestamp is resolved on the fly to + the earliest decoded frame at or after it — never an earlier one, so a + camera observation can't predate the action paired with it (the same + selection rule as ``FrameCache.get_frame_at``). Targets past the end of + the stream clamp to the last decoded frame. Only selected frames are + retained, already resized to ``target_size``, so peak memory is one raw + decoded frame plus the selected output frames. Decoding stops early + once every target is resolved. Args: packets: Packet series from :meth:`load_packets`. @@ -432,42 +431,36 @@ def decode_at( num_targets = len(targets) frames_out: list[npt.NDArray[np.uint8]] = [] - prev_frame: npt.NDArray[np.uint8] | None = None - prev_ts = 0 - prev_out: npt.NDArray[np.uint8] | None = None # processed prev_frame + cur_frame: npt.NDArray[np.uint8] | None = None + cur_ts = 0 + cur_out: npt.NDArray[np.uint8] | None = None # processed cur_frame - def emit_prev() -> npt.NDArray[np.uint8]: + def emit_cur() -> npt.NDArray[np.uint8]: # Resize lazily and once per selected source frame. - nonlocal prev_out - if prev_out is None: - assert prev_frame is not None + nonlocal cur_out + if cur_out is None: + assert cur_frame is not None if target_size is not None: - prev_out = resize_rgb(prev_frame, *target_size) + cur_out = resize_rgb(cur_frame, *target_size) else: - prev_out = prev_frame - return prev_out + cur_out = cur_frame + return cur_out decoded_frames = 0 for frame, ts in self._iter_frames(packets): decoded_frames += 1 - if prev_frame is None: - prev_frame, prev_ts = frame, ts - continue + cur_frame, cur_ts, cur_out = frame, ts, None - # All targets at or before the midpoint of (prev, current) are - # nearest to prev (ties go to the earlier frame). - while len(frames_out) < num_targets and ( - targets[len(frames_out)] - prev_ts - ) <= (ts - targets[len(frames_out)]): - frames_out.append(emit_prev()) + # This is the earliest decoded frame at or after every remaining + # target up to `ts` — emit it for all of them. + while len(frames_out) < num_targets and targets[len(frames_out)] <= ts: + frames_out.append(emit_cur()) if len(frames_out) >= num_targets: # Every target resolved — skip decoding the rest of the stream. break - prev_frame, prev_ts, prev_out = frame, ts, None - - if prev_frame is None: + if cur_frame is None: logger.warning( "No frames decoded from {} packets for {}", packets.num_packets, @@ -475,9 +468,20 @@ def emit_prev() -> npt.NDArray[np.uint8]: ) return FrameCache() - # Remaining targets are at/after the last frame: clamp to it. + if len(frames_out) < num_targets: + # Targets past the last decoded frame clamp to it — the one case + # where the at-or-after rule can still under-shoot the target. + logger.warning( + "{} of {} targets for {} fall past the last decoded frame " + "(last frame {}ns, last target {}ns) — clamped to it", + num_targets - len(frames_out), + num_targets, + packets.entity, + cur_ts, + int(targets[-1]), + ) while len(frames_out) < num_targets: - frames_out.append(emit_prev()) + frames_out.append(emit_cur()) logger.info( "Decoded {} frames from {} packets for {}, kept {} grid frames", diff --git a/tests/test_export.py b/tests/test_export.py index 47715c0..ecb8633 100644 --- a/tests/test_export.py +++ b/tests/test_export.py @@ -223,8 +223,8 @@ def test_single_frame(self): result = cache.get_frame_at(5000) assert np.array_equal(result, frame) - def test_nearest_frame_lookup(self): - """Lookup returns nearest frame by timestamp.""" + def test_at_or_after_frame_lookup(self): + """Lookup returns the earliest frame at or after the timestamp.""" frames = [create_test_frame(value=i * 50) for i in range(5)] timestamps = np.array([0, 1000, 2000, 3000, 4000], dtype=np.int64) cache = FrameCache(frames=frames, timestamps_ns=timestamps) @@ -239,14 +239,12 @@ def test_nearest_frame_lookup(self): assert np.array_equal(cache.get_frame_at(2000), frames[2]) assert np.array_equal(cache.get_frame_at(4000), frames[4]) - # Nearest to lower - assert np.array_equal(cache.get_frame_at(400), frames[0]) - - # Nearest to upper + # Never earlier than the target: all of these round *up* to frame 1 + assert np.array_equal(cache.get_frame_at(1), frames[1]) + assert np.array_equal(cache.get_frame_at(400), frames[1]) + assert np.array_equal(cache.get_frame_at(500), frames[1]) assert np.array_equal(cache.get_frame_at(600), frames[1]) - - # Midpoint goes to lower - assert np.array_equal(cache.get_frame_at(500), frames[0]) + assert np.array_equal(cache.get_frame_at(999), frames[1]) # Before start assert np.array_equal(cache.get_frame_at(-1000), frames[0]) @@ -264,13 +262,19 @@ def test_frame_index_lookup(self): assert cache.get_frame_index(200) == 1 assert cache.get_frame_index(300) == 2 - # Nearest (equidistant rounds to earlier frame with <=) - assert cache.get_frame_index(140) == 0 + # Non-exact targets round *up* — never to a frame before the target + assert cache.get_frame_index(101) == 1 + assert cache.get_frame_index(140) == 1 assert cache.get_frame_index(160) == 1 - assert ( - cache.get_frame_index(250) == 1 - ) # Equidistant (50 from 200 and 300) -> earlier - assert cache.get_frame_index(251) == 2 # Closer to 300 + assert cache.get_frame_index(250) == 2 + assert cache.get_frame_index(251) == 2 + + # Before the first frame -> first frame; past the last -> last frame + assert cache.get_frame_index(50) == 0 + assert cache.get_frame_index(1000) == 2 + + # Empty cache + assert FrameCache().get_frame_index(100) == -1 # =============================================================================