scripts: copy paired annotations between zarr recordings

[ElmoPA]

No description provided.

[ElmoPA]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• uv.lock: regenerate #398
• scripts: copy paired annotations between zarr recordings #397 👈 (View in Graphite)
• scripts: latent inspector, velocity distribution, zarr validator #396
• eval: latent-feature evaluator and cotrain config #395
• rldb: SafeS3EpisodeResolver + EvenStrideDataset; harden cotrain configs #394
• rldb: surface episode_hash on every sample #393
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[graphite-app]

Silent propagation of corrupted data. If the source zarr array exists but cannot be read (e.g., corrupted), _entries(src) returns -1. The script continues copying, then _entries(dst) also returns -1, and the verification at line 123 passes (-1 == -1), falsely indicating success.

Fix by checking for valid source data:

src_n = _entries(src)
if src_n < 0:
    raise SystemExit(f"Source exists but cannot be read as zarr: {src}")

Suggested change

	src_n = _entries(src)
	dst_n_before = _entries(dst) if os.path.exists(dst) else 0
	print(
	f"[{'dry' if args.dry_run else 'copy'}] {args.name}: "
	f"{src_hash}({src_n}) -> {dst_hash} (had {dst_n_before})"
	)
	if args.dry_run:
	continue
	if os.path.exists(dst):
	shutil.rmtree(dst)
	shutil.copytree(src, dst)
	verify_n = _entries(dst)
	print(f" verified: n_entries={verify_n}")
	if verify_n != src_n:
	raise SystemExit(f"Mismatch after copy: src={src_n} dst={verify_n}")
	src_n = _entries(src)
	if src_n < 0:
	raise SystemExit(f"Source exists but cannot be read as zarr: {src}")
	dst_n_before = _entries(dst) if os.path.exists(dst) else 0
	print(
	f"[{'dry' if args.dry_run else 'copy'}] {args.name}: "
	f"{src_hash}({src_n}) -> {dst_hash} (had {dst_n_before})"
	)
	if args.dry_run:
	continue
	if os.path.exists(dst):
	shutil.rmtree(dst)
	shutil.copytree(src, dst)
	verify_n = _entries(dst)
	print(f" verified: n_entries={verify_n}")
	if verify_n != src_n:
	raise SystemExit(f"Mismatch after copy: src={src_n} dst={verify_n}")

Spotted by Graphite



Is this helpful? React 👍 or 👎 to let us know.

[github-actions]

Claude Code Review

Summary

Adds a one-off utility script to copy a named zarr sub-array (default annotations) from one episode directory to another via shutil.copytree. Useful for propagating language annotations from robot to paired human episodes.

Key concerns

1. Bypasses ZarrWriter convention. The codebase rule is "Always use ZarrWriter, never write zarr manually." shutil.copytree is a filesystem-level copy that sidesteps any consolidated metadata, manifest, or sibling-array consistency the writer would maintain. For a Zarr v3 store, the parent group's zarr.json likely needs to acknowledge the new array, and other arrays in the destination episode (length-aligned) may now be inconsistent with the copied one. At minimum this should:

   • Open the destination episode group via ZarrWriter (or zarr.open_group) and write the array through it.
   • Re-consolidate metadata after copy.

2. No length-alignment check against the destination. The script verifies src_n == dst_n_after_copy (trivially true for a copytree), but does not check that src_n matches the destination episode's other per-step arrays (e.g., actions, obs/*). Pasting an annotations array of length N into an episode of length M will silently corrupt training-time indexing.

3. Default --root is a hardcoded user path (/storage/project/r-dxu345-0/agao81/pick_place). For a script that is destructive (shutil.rmtree(dst)) this should have no default, or at least no user-specific default committed to the repo.

4. Cross-embodiment copy is unsafe by default. The motivating use case is robot→human pairing, i.e. different embodiments. Copying annotations is probably fine, but the same script with --name dino.front_1 or --name images.front_1 (explicitly suggested in the docstring) would copy embodiment-specific sensor data across embodiments, which is almost certainly wrong. Consider whitelisting --name to known-safe arrays, or at least loudly warning when --name is not annotations.

5. No tests. Even a small pytest with a tmp zarr store covering: missing src, mismatched lengths, dst-overwrite, and metadata consolidation would catch the issues above.

Suggestions

• Replace shutil.copytree with a proper zarr read → write through ZarrWriter (or zarr.copy / zarr.copy_store) into the destination episode group, then consolidate metadata.
• After write, validate that len(dst_array) == len(dst_episode["actions"]) (or whichever canonical per-step array the episode uses) and abort otherwise.
• Drop the hardcoded --root default; make it required.
• Either restrict --name to annotations or gate other names behind an explicit --allow-unsafe-name flag.
• Log episode hashes that you're about to rmtree and consider a confirmation prompt or --force.
• Add a small unit test under tests/ exercising the happy path and the length-mismatch failure.

Verdict: Request Changes

The script is small and the intent is clear, but writing into zarr stores via shutil violates a stated codebase invariant and there's no guard against length-mismatched annotations silently entering training data. Worth a small refactor before merge.

---

Reviewed by Claude · Review workflow