Fix bug in dataset root resolution

This commit is contained in:
Saberlve 2026-08-08 16:51:08 +08:00
parent ed0063730d
commit 1b775a4be2
4 changed files with 56 additions and 60 deletions

View File

@ -449,9 +449,9 @@ class UFRobot(Robot, Thread):
warn_code = getattr(arm, "warn_code", "unknown") warn_code = getattr(arm, "warn_code", "unknown")
return f"mode={mode}, state={state}, error_code={error_code}, warn_code={warn_code}" return f"mode={mode}, state={state}, error_code={error_code}, warn_code={warn_code}"
def _check_motion_code(self, command: str, code: int) -> None: def _check_motion_code(self, command: str, code: int | None) -> None:
"""Fail loudly when the SDK rejects a joint command.""" """Fail loudly when the SDK explicitly rejects a motion command."""
if code != 0: if code is not None and code != 0:
raise RuntimeError(f"{command} failed, code={code}, {self._motion_status()}") raise RuntimeError(f"{command} failed, code={code}, {self._motion_status()}")
def send_action(self, action: dict) -> np.ndarray: def send_action(self, action: dict) -> np.ndarray:

View File

@ -510,26 +510,40 @@ def _ask_choice(prompt: str, options: dict[str, str]) -> str:
print(f"Invalid choice, please enter {keys}.") print(f"Invalid choice, please enter {keys}.")
def _missing_dataset_files(root: Path) -> list[str]:
"""Return the local files required before a dataset can be resumed."""
missing = []
for relative_path in ("meta/info.json", "meta/tasks.parquet"):
if not (root / relative_path).is_file():
missing.append(relative_path)
if not any((root / "meta" / "episodes").glob("*/*.parquet")):
missing.append("meta/episodes/*/*.parquet")
if not any((root / "data").glob("*/*.parquet")):
missing.append("data/*/*.parquet")
return missing
def _prepare_dataset_root(cfg: UFRecordConfig) -> None: def _prepare_dataset_root(cfg: UFRecordConfig) -> None:
"""Create the dataset root and ask how to handle an existing dataset.""" """Prepare an existing dataset root without pre-creating a new one."""
root = Path(cfg.dataset.root) root = Path(cfg.dataset.root)
existed = root.exists() existed = root.exists()
# Create the dataset root (and any parent directories) first. if not existed:
root.mkdir(parents=True, exist_ok=True) if cfg.resume:
raise RuntimeError(f"Cannot resume because the dataset directory does not exist: {root}")
if not existed or cfg.resume:
return return
if not (root / "meta" / "info.json").is_file(): missing = _missing_dataset_files(root)
# The directory exists but is not a valid LeRobot dataset. if missing:
if not sys.stdin.isatty(): missing_text = ", ".join(missing)
raise RuntimeError( message = (
f"Dataset directory exists but is not a valid LeRobot dataset: {root}\n" f"Dataset directory is incomplete and cannot be resumed: {root}\n"
"Choose a new dataset.root, or remove this empty/incomplete directory before recording." f"Missing: {missing_text}"
) )
if cfg.resume or not sys.stdin.isatty():
raise RuntimeError(message)
choice = _ask_choice( choice = _ask_choice(
f"Directory exists but is not a valid LeRobot dataset: {root}", message,
options={ options={
"o": "Overwrite: remove this directory and record a new dataset", "o": "Overwrite: remove this directory and record a new dataset",
"c": "Cancel", "c": "Cancel",
@ -541,6 +555,9 @@ def _prepare_dataset_root(cfg: UFRecordConfig) -> None:
raise SystemExit("Recording cancelled.") raise SystemExit("Recording cancelled.")
return return
if cfg.resume:
return
# A valid LeRobot dataset already exists. # A valid LeRobot dataset already exists.
if not sys.stdin.isatty(): if not sys.stdin.isatty():
# Non-interactive run: keep the previous auto-resume behaviour. # Non-interactive run: keep the previous auto-resume behaviour.

View File

@ -2,49 +2,7 @@
set -euo pipefail set -euo pipefail
repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
cd "$repo_root"
config_path="config/manual_mode/xarm7_manual_record_config.yaml"
dataset_root="$(
sed -n '/^dataset:/,/^[^[:space:]]/s/^[[:space:]]*root:[[:space:]]*//p' "$config_path" \
| head -n 1 \
| tr -d "\"'"
)"
if [[ -z "$dataset_root" ]]; then
printf 'Could not read dataset.root from %s\n' "$config_path" >&2
exit 1
fi
if [[ "$dataset_root" != /* ]]; then
dataset_root="$repo_root/$dataset_root"
fi
record_args=("$@")
resume_requested=false
for arg in "${record_args[@]}"; do
if [[ "$arg" == "-r" ]]; then
resume_requested=true
break
fi
done
if [[ -e "$dataset_root" ]]; then
if [[ ! -f "$dataset_root/meta/info.json" ]]; then
printf 'Dataset directory exists but is not a valid LeRobot dataset: %s\n' "$dataset_root" >&2
printf 'Choose a new dataset.root, or remove this empty/incomplete directory before recording.\n' >&2
exit 1
fi
if [[ "$resume_requested" == false ]]; then
record_args=("-r" "${record_args[@]}")
fi
elif [[ "$resume_requested" == true ]]; then
printf 'Cannot resume because the dataset directory does not exist: %s\n' "$dataset_root" >&2
exit 1
fi
exec uv run uf-lerobot-record \ exec uv run uf-lerobot-record \
--config_path "$config_path" \ --config_path config/manual_mode/xarm7_manual_record_config.yaml \
"${record_args[@]}" "$@"

View File

@ -315,6 +315,27 @@ def test_prepare_dataset_root_rejects_incomplete_resume(tmp_path):
_prepare_dataset_root(cfg) _prepare_dataset_root(cfg)
def test_prepare_dataset_root_rejects_resume_when_root_is_missing(tmp_path):
cfg = SimpleNamespace(dataset=SimpleNamespace(root=tmp_path / "missing"), resume=True)
with pytest.raises(RuntimeError, match="does not exist"):
_prepare_dataset_root(cfg)
def test_prepare_dataset_root_resumes_complete_dataset_without_prompt(tmp_path, monkeypatch):
root = tmp_path / "dataset"
(root / "meta" / "episodes" / "chunk-000").mkdir(parents=True)
(root / "data" / "chunk-000").mkdir(parents=True)
(root / "meta" / "info.json").write_text("{}")
(root / "meta" / "tasks.parquet").write_bytes(b"tasks")
(root / "meta" / "episodes" / "chunk-000" / "file-000.parquet").write_bytes(b"episodes")
(root / "data" / "chunk-000" / "file-000.parquet").write_bytes(b"data")
cfg = SimpleNamespace(dataset=SimpleNamespace(root=root), resume=True)
monkeypatch.setattr(record_module.sys.stdin, "isatty", lambda: True)
_prepare_dataset_root(cfg)
def test_manual_record_loop_writes_actual_state_as_action(tmp_path): def test_manual_record_loop_writes_actual_state_as_action(tmp_path):
class FakeRobot: class FakeRobot:
name = "fake_manual_robot" name = "fake_manual_robot"