Fast Openarm dataset to Lerobot dataset conversion - #58
Open
k1000dai wants to merge 13 commits into
Open
Conversation
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Load+downsample, image-stats decode, and ffmpeg video encoding now run in parallel across episodes (processes for CPU/IO work, threads for ffmpeg subprocesses), keeping assembly serial and ordered so output is identical to the serial path. Adds a --jobs/-j flag (default: all cores; 1 = serial). v3.0 packed-video file boundaries are now planned up front from a fixed per-camera compression ratio (instead of an online feedback loop) so every file encodes independently in parallel; sizes still track the target approximately and single-pass-per-file (uniform PTS) is preserved. Warm benchmark on the 100-episode sample (256-core node): v2.1 ~797s -> 88s (9.0x), v3.0 ~630s -> 131s (4.8x). Output verified byte/numeric-identical to the previous serial code for both formats. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Each shard converts a contiguous slice of episodes into a self-contained partial LeRobot dataset (global episode/task indices assigned at shard time, local row-index/file numbering). aggregate_shards stitches shards by moving/renumbering data+video files and fixing up the global row index -- without re-encoding video. Supports lerobot_v2.1, lerobot_v3.0, and gr00t. New: - distributed.py: convert_shard, aggregate_shards (+ helpers) - aggregate.py, slurm_submit.py CLIs; convert.py gains --num-shards/--shard-index - openarm-dataset-aggregate / openarm-dataset-slurm-submit entry points - to_lerobotv21/to_lerobotv30 refactored to share _write_v21/_write_v30 bodies Sharded+aggregated output is verified equal to a single-process conversion (per-episode data, stats.json, info totals) for all three formats; video file layout differs (sharded packing) but is valid. Tests need no cluster/lerobot. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Previously a failing sbatch raised an opaque CalledProcessError with the real message hidden by capture_output. Now the sbatch stderr (e.g. 'Requested node configuration is not available') is raised in a clear RuntimeError with a hint. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
parallel_map/thread_map gain a desc that advances a tqdm bar as items complete (order and fail-fast preserved via submit + as_completed). Wired into the load-episodes and video-encode stages of v2.1 and v3.0 so long runs -- and SLURM shard logs -- show live progress. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
resolve_jobs and ffmpeg_threads_for now use sched_getaffinity (available_cpus) instead of os.cpu_count(), so a --cpus-per-task allocation is honored: ffmpeg no longer spawns threads for all physical cores inside a smaller cgroup, and the default (all cores) means all *allocated* cores. slurm-submit gains --jobs to set the in-shard worker count independently of --cpus-per-task (loading is NFS-latency bound; matching/exceeding cpus is best). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The SLURM shard->aggregate path (Phase B) added a large aggregation layer (distributed.py plus aggregate/slurm-submit CLIs and generalized shared writers) and a cluster dependency, only winning when one node's cores are the bottleneck and multiple nodes are available. For this workload ffmpeg saturates a single node, so in-process parallelism (Phase A) is simpler, dependency-free, and effectively faster. Keep only Phase A. Removed: - distributed.py, aggregate.py, slurm_submit.py, tests/test_distributed.py - openarm-dataset-aggregate / openarm-dataset-slurm-submit entry points - convert.py --num-shards / --shard-index and the shard branch - Phase B design spec Reverted the Phase-B generalizations in the shared writers back to Phase A: - _collect_downsampled_data drops the shard-only episode_indices param - _write_v21 / _write_v30 inlined back into to_lerobotv21 / to_lerobotv30 Sharding work is preserved on the `sharding` branch. Full suite green (pytest), ruff/shfmt clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds a parallel execution path for converting OpenArm datasets into LeRobot v2.1 / v3.0 (and GR00T) outputs, exposing a --jobs/-j CLI flag and introducing shared parallel utilities to speed up episode processing and video encoding while keeping serial/parallel outputs identical (except for MP4 bytes).
Changes:
- Add
src/openarm_dataset/parallel.pyhelpers (parallel_map,thread_map, CPU detection, ffmpeg thread budgeting). - Parallelize episode loading/downsampling + image-stat computation (process pool) and video encoding (thread pool) for LeRobot v2.1/v3.0 conversions; add
jobsplumbing through the converters and CLI. - Add regression tests ensuring parallel conversion matches serial for parquet/meta outputs and that videos are non-empty.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_lerobot_parallel.py | Adds parity tests for serial vs parallel conversion and default-jobs execution. |
| src/openarm_dataset/parallel.py | Introduces reusable process/thread mapping helpers and CPU/ffmpeg thread resolution. |
| src/openarm_dataset/lerobot_v30.py | Refactors v3.0 video packing/encoding and stats flow to support parallel execution with jobs. |
| src/openarm_dataset/lerobot_v21.py | Parallelizes per-episode processing and video encoding for v2.1/GR00T; adds jobs parameter. |
| src/openarm_dataset/ffmpeg.py | Adds optional threads argument and avoids Path.resolve() syscalls in concat file path handling. |
| src/openarm_dataset/convert.py | Adds --jobs/-j CLI flag and passes it to LeRobot/GR00T conversions. |
Comment on lines
+155
to
+156
| step = max(1, len(non_empty) // max_samples) | ||
| sampled_frames = [fl[len(fl) // 2] for fl in non_empty[::step]] |
Comment on lines
+272
to
+282
| # Calibrate a compression ratio per camera (one sample encode each), in | ||
| # parallel across cameras. | ||
| ratios = thread_map( | ||
| lambda name: _calibrate_compression_ratio( | ||
| cam_frames[name][0] if cam_frames[name] else [], | ||
| cam_src_sizes[name][0] if cam_src_sizes[name] else 0.0, | ||
| fps, | ||
| ), | ||
| image_names, | ||
| jobs, | ||
| ) |
Comment on lines
+284
to
+286
| # Plan file boundaries per camera and collect every file's encode job. | ||
| encode_jobs: list[tuple[Path, list]] = [] | ||
| for name, ratio in zip(image_names, ratios): |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.