Repository navigation
[python] Make LeRobot Dataset appends idempotent across Ray jobs - #11097
XiaoHongbo-Hope wants to merge 2 commits into
Conversation
JingsongLi
left a comment
There was a problem hiding this comment.
Reviewed commit 8494d416a1. Three inline findings below, covering coordinator resource allocation, table-group identity, and compatibility with existing strict-schema append workflows.
| name = "lerobot_append_" + hashlib.sha256(identity.encode("utf-8")).hexdigest() | ||
| actor = ray.remote(_AppendCoordinator).options( | ||
| name=name, namespace=_NAMESPACE, lifetime="detached", | ||
| get_if_exists=True, max_concurrency=1, num_cpus=1, |
There was a problem hiding this comment.
[P1] Avoid reserving a CPU for the coordinator's lifetime
This detached actor keeps its explicit num_cpus=1 allocation even while idle or waiting for preparation tasks. Those tasks also require one CPU in _prepare_sources_ray, so a coordinated append with parallelism > 1 cannot make progress on a one-CPU cluster. On a larger cluster, coordinators for enough groups can consume every CPU and block further preparation or actor creation. I reproduced the scheduling deadlock with these actor options on Ray 2.59.1; changing only the coordinator to num_cpus=0 allowed the child task to finish. Please avoid retaining a CPU for the coordinator's lifetime, or otherwise manage its resources so preparation can always run.
| frames_name = _companion_identifier( | ||
| connection._identifier(table_name), "_frames") | ||
| frames = connection.get_table(frames_name).raw_table | ||
| return frames.catalog_environment.uuid or frames.table_path |
There was a problem hiding this comment.
[P1] Use a canonical identity for filesystem table groups
For filesystem catalogs, catalog_environment.uuid is None and table_path retains the warehouse's original spelling. Two jobs using /shared/warehouse and file:///shared/warehouse therefore generate different actor names for the same physical table group, even though both use coordinated=True on the same Ray cluster. I verified this with a real LeRobot group: both connections read and write the same tables, but _group_identity differs by the file:// prefix. The two actors can then run append/recovery concurrently, reintroducing the race this PR is intended to prevent. Please use a persisted table/group ID or canonicalize the storage identity before deriving the actor name.
| _DATASETS_SCHEMA = pa.schema([ | ||
| pa.field("index", pa.int64(), nullable=False), | ||
| pa.field("name", pa.string()), | ||
| pa.field("request_id", pa.string()), |
There was a problem hiding this comment.
[P2] Preserve ordinary strict-schema appends to existing groups
These columns are now included in every incoming Dataset schema, including direct appends with request_id=None. A table group created before this change lacks them, so appending an otherwise compatible source with schema_evolution=False now raises LeRobot source 1 adds field request_id with schema_evolution disabled, even when coordination is not enabled and no business columns have changed. I reproduced the full import/append path with the previous Dataset schema; the same append succeeds with the baseline schema. Please avoid requiring request metadata for ordinary legacy appends, or handle internal metadata migration separately from feature-schema evolution.
|
Still not very clear now. still designing @JingsongLi |
JingsongLi
left a comment
There was a problem hiding this comment.
Reviewed commit e958b89873. Two compatibility issues were reproduced locally; details are in the inline comments.
| has_coordinated_rows = any(row.get("dataset_id") is not None | ||
| for row in source_rows) |
There was a problem hiding this comment.
[P2] Distinguish ordinary idempotent appends from coordinated writes
A non-null dataset_id does not imply coordinated=True: the API and new tests explicitly support supplying an ID for an ordinary single-writer append, and _plan_merge saves that ID in _datasets in either mode. I reproduced an initial list import followed by append_to_lerobot(..., dataset_id="b"); that append succeeds, but the next ordinary append, with either a new ID or no ID, raises LeRobot group has coordinated Datasets. This prevents continuing the supported single-writer workflow after its first idempotent append. Please detect the actual coordinated storage mode, for example through an explicit marker or its physical-key allocation, instead of treating every saved ID as evidence of coordination.
| for name in ("dataset_id", "source_hash"): | ||
| if existing is None or coordinated or name in _target_schema(table.raw_table).names: | ||
| target_schema = target_schema.append(pa.field(name, pa.string())) |
There was a problem hiding this comment.
[P2] Preserve custom features that share the new identity-column names
For a new group, these identity columns are appended even when no dataset_id is supplied and coordinated=False. A previously valid source feature named dataset_id or source_hash therefore produces duplicate schema fields and makes an ordinary import fail. Using the existing test helpers, I added an int64 dataset_id feature and called load_from_lerobot("robot", [uri]): base 40a386e12f imports and preserves the values, while this commit raises ValueError ... must not contain duplicate fields. Found: ['dataset_id']. Please use collision-safe internal identity storage or otherwise preserve ordinary imports whose source schema already owns these names.
JingsongLi
left a comment
There was a problem hiding this comment.
Re-reviewed commit e395af6160. The previous two compatibility findings are fixed. Two additional issues were reproduced locally and are described inline.
| if not commit_result_may_be_uncertain and prepared_messages is not None: | ||
| self.abort(prepared_messages) |
There was a problem hiding this comment.
[P1] Preserve prepared files across uncertain commit retries
commit_result_may_be_uncertain is reset for each _try_commit call, so this check cannot establish that the messages were never published by an earlier call. I reproduced this through a conditional StreamTableCommit: with commit.max-retries=0, publish the snapshot and then raise an OSError to simulate a lost response. The first call raises while snapshot 1 and its data file exist. Retrying the same checkpoint with the same prepared messages makes the condition return False because the append is already present, but the fresh local uncertainty flag permits abort, deleting the file referenced by snapshot 1; reads then fail with FileNotFoundError. Please retain or resolve uncertain ownership across calls, or verify that the prepared files are not referenced before cleaning them up. The current invocation's failure to publish does not prove that these files are uncommitted.
| key_column = ("dataset_index" if "dataset_index" in | ||
| _target_schema(pinned).names else "index") |
There was a problem hiding this comment.
[P2] Select the Dataset key explicitly instead of inferring its table role
extra_dataset_fields={"dataset_index": "notes.txt"} is accepted and creates a user string column on _datasets. Its presence makes this condition treat the Dataset table as a Frames/Episodes table and filter the text column instead of the physical key index. For ordinary notes text, every check reports the Dataset absent. I reproduced two concurrent first attempts with the same dataset_id, synchronizing them after Frame preparation: both succeeded, leaving 10 Frames, 4 Episodes, and 3 Dataset rows, including two identical physical Dataset keys. The training reader then fails with index must be unique and non-negative. The early replay check cannot prevent this race because neither attempt has published metadata yet. Please pass the table role or key/identity columns explicitly; Dataset publication must always check index, regardless of custom metadata columns.
02b9317 to
caa49a9
Compare
JingsongLi
left a comment
There was a problem hiding this comment.
Requirement fit: SUPPORTED. Implementation: FINDINGS.
Independent worktree review of the current UUID-based implementation, including conditional commit ownership, recovery, Dataset-last publication and training readers. The previous role-selection, namespaced-column and reused-stream-condition fixes were rechecked. Two new failures were reproduced with actual filesystem tables: a conditional retry deletes already-published files after a callback error, and a first coordinated append can leak unpublished Episode statistics into a reader. Details inline.
Validation: 482 selected existing commit/LeRobot/Ray/writer/text-column tests passed across runs, with 2 platform/dependency skips. Initial LeRobot fixture errors were resolved with compatible NumPy in an isolated environment. Some completed pytest runs stalled in interpreter finalization waiting for child processes and were terminated after their test summaries; this is not claimed as a clean process exit. Both focused failure probes still fail on this head. Project-configured flake8 and diff whitespace checks passed. No production source modifications were made.
| if not needed: | ||
| # An uncertain response may have published these files. | ||
| if not commit_result_may_be_uncertain and prepared_messages is not None: | ||
| self.abort(prepared_messages) |
There was a problem hiding this comment.
[P1] Preserve published files when a commit callback fails
A snapshot can already be published while commit_outcome_uncertain remains false: _try_commit_once calls _notify_commit_callbacks after the publication try/except, so a callback exception escapes without marking the prepared messages. Retrying the same conditional stream commit then sees its condition become false and aborts these files. I reproduced this through the public API: prepare one row, install a CommitCallback that raises OSError, use condition snapshot is None, commit identifier 1, then retry the same identifier/messages. Snapshot 1 exists after the first exception; the retry returns false and deletes the data file referenced by snapshot 1. Reads are subsequently broken. Record published ownership before invoking callbacks, or recognize this commit's publication before aborting on a false condition; checking only publication-response uncertainty is insufficient.
| and episode_snapshot.next_row_id != len(episodes)) | ||
| stats = _merged_source_stats( | ||
| source_rows, episodes_table, | ||
| frames if recompute_controls else None, |
There was a problem hiding this comment.
[P2] Keep Episode statistics on the checked snapshot
During the first coordinated append to an ordinary imported group, the reader can observe the old Episode snapshot here and choose recompute_controls=False, then _merged_source_stats scans a newer, unpinned Episode snapshot. That fallback does not filter statistics by the captured published Dataset IDs. A deterministic interleaving that appends Frames/Episodes after this check and fails before Dataset publication gives a successful reader with 5 published frames but meta.stats['index']['count'] == [12]; control moments include the unpublished 7-frame source as well. No published coordinated marker exists yet to force the safe path. Pin/reuse the checked Episode snapshot, or filter/rebuild statistics strictly from the captured published sources.
b2a1209 to
6486885
Compare
6486885 to
c59b319
Compare
Independent Ray jobs can append LeRobot sources to one three-table group by supplying
coordinated=Trueand a stabledataset_id. The Dataset identity is a UUID in__datasets.id, referenced by__frames.dataset_idand__episodes.dataset_id. A caller UUID is used directly; another stable caller key is converted to UUIDv5. Ordinary imports continue to use UUIDv7.Commit conditions deduplicate concurrent writes of the same source. Frames and Episodes can be resumed independently after a failure, and the Dataset row is published last. Readers ignore components without a published Dataset row and rebuild training control statistics from source-local counts in current UUID order. A reused identity with different import arguments is rejected. The request fingerprint does not hash source bytes, so sources must stay immutable across retries.
The commit layer cleans prepared files for definite condition failures and preserves them when a snapshot publication result may be uncertain, including across a retry with serialized messages. The Dataset condition selects its key by table role, so an optional user text field named
dataset_indexcannot affect deduplication. It checks duplicate Frames per split and batch, without materializing the full Dataset in Python.Reusable stream commits propagate a newly set condition to an existing Python backend. Internal Dataset metadata uses
__paimon_source_hash,__paimon_coordinated, and__paimon_control_counts, leaving the unprefixed names available for historical user text columns.A successful snapshot now marks prepared messages before invoking commit callbacks, so a callback exception followed by a conditional retry cannot delete published files. When an older table group lacks stored control histograms, its first coordinated append persists them in the Dataset table; later reader initialization does not rescan its Frames. Concurrent backfills use a snapshot condition, and the losing writer reuses the published counts.
Count-based
recover_lerobot_appendrejects groups with published coordinated Datasets because later publications make row-count rollback unsafe; retry an interrupted coordinated append with the samedataset_id. A coordinated append rejectstag_name; create a tag separately after publication withcreate_lerobot_tag.Reader statistics now filter Episode rows to the captured published Dataset IDs, so an Episode commit racing with reader initialization cannot contribute unpublished control stats. Validation after rebasing onto master
6d50c1ab08: 444 LeRobot import, reader, writer, and commit tests passed; 1 skipped. Flake8 on changed Python files andgit diff --checkpassed.