JingsongLi commented on code in PR #10483: URL: https://github.com/apache/paimon/pull/10483#discussion_r4242063875
########## paimon-python/pypaimon/multimodal/lerobot/ray_loader.py: ########## @@ -0,0 +1,166 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ray workers prepare LeRobot files; only the driver commits snapshots.""" + +import logging +from pathlib import Path + +from pypaimon.multimodal.lerobot.loader import ( + _abort_messages, _commit_messages, _prepare_dataset, +) +from pypaimon.multimodal.lerobot.source import ( + _RemoteLeRobotDataset, _resolved_source, +) + + +_CANCEL_DRAIN_SECONDS = 1.0 + + +def _require_ray(): + try: + import ray + except ImportError as error: + raise ImportError("Parallel import requires 'pypaimon[ray,lerobot]'.") from error + if not ray.is_initialized(): + raise RuntimeError("Call ray.init() before using parallelism > 1.") + return ray + + +def _ray_source_uri(source): + if isinstance(source, str) and "://" in source: + return source + if isinstance(source, Path) or isinstance(source, str) and Path(source).is_dir(): + return Path(source).expanduser().resolve().as_uri() Review Comment: [P2] Preserve absolute image paths within a local source directory Converting a shared local directory to `file://` switches the import to `_RemoteLeRobotDataset`. An image represented as `{"bytes": None, "path": "/shared/source/images/0.png"}` is supported by the serial path: `_image_bytes` checks that the resolved path remains inside the source and reads it. The parallel reader instead rejects this unschemed absolute path in `_remote_source_path`, so enabling parallelism makes the same valid dataset fail. I reproduced this with the same fixture and real Paimon writer; `test_dataset_paths_cannot_escape_source` also explicitly covers the supported serial behavior. Please preserve that containment-checked local path handling, or normalize these paths to file URIs, and add a regression case with path-backed images. ########## paimon-python/pypaimon/multimodal/lerobot/ray_loader.py: ########## @@ -0,0 +1,166 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ray workers prepare LeRobot files; only the driver commits snapshots.""" + +import logging +from pathlib import Path + +from pypaimon.multimodal.lerobot.loader import ( + _abort_messages, _commit_messages, _prepare_dataset, +) +from pypaimon.multimodal.lerobot.source import ( + _RemoteLeRobotDataset, _resolved_source, +) + + +_CANCEL_DRAIN_SECONDS = 1.0 + + +def _require_ray(): + try: + import ray + except ImportError as error: + raise ImportError("Parallel import requires 'pypaimon[ray,lerobot]'.") from error + if not ray.is_initialized(): + raise RuntimeError("Call ray.init() before using parallelism > 1.") + return ray + + +def _ray_source_uri(source): + if isinstance(source, str) and "://" in source: + return source + if isinstance(source, Path) or isinstance(source, str) and Path(source).is_dir(): + return Path(source).expanduser().resolve().as_uri() + raise ValueError( + "Parallel import requires a FileIO URI or shared directory, not a Hub repo_id.") + + +def _episode_groups(info, episodes, video_fields): + # Keep Episodes sharing a data or video file in the same task. + parents = list(range(len(episodes))) + + def root(i): + while parents[i] != i: + parents[i] = parents[parents[i]] + i = parents[i] + return i + + owners = {} + for i in range(len(episodes)): + episode = episodes[i] + data_path = _RemoteLeRobotDataset._data_path(episode, info) + previous = owners.setdefault(("data", data_path), i) + parents[root(i)] = root(previous) + for field in video_fields: + prefix = "videos/%s/" % field + path = info["video_path"].format( + video_key=field, chunk_index=episode[prefix + "chunk_index"], + file_index=episode[prefix + "file_index"]) + previous = owners.setdefault(("video", path), i) + parents[root(i)] = root(previous) + groups = {} + for i in range(len(episodes)): + groups.setdefault(root(i), []).append(i) + return list(groups.values()) + + +def _task_groups(info, episodes, video_fields, parallelism): + groups = _episode_groups(info, episodes, video_fields) + tasks = [[] for _ in range(min(parallelism, len(groups)))] + for i, group in enumerate(groups): + tasks[i % len(tasks)].extend(group) + return [sorted(task) for task in tasks] + + +def _prepare_group(table, source_path, source_options, info, schema, + batch_size, metadata, video_fields, episode_indices): + with _resolved_source(source_path, source_options) as (source, current_info): + if current_info != info: + raise ValueError("LeRobot source metadata changed during import.") + dataset = _RemoteLeRobotDataset(source, info) + messages = None + close_started = False + try: + messages = _prepare_dataset( + table, dataset, info, source, schema, batch_size, + metadata, video_fields, episode_indices) + close_started = True + dataset.close() + return messages + except BaseException: + # _prepare_dataset aborts an unfinished writer. A cancellation + # after prepare_commit must also discard its returned files. + if messages is not None: + try: + _abort_messages(table, messages) + except Exception: + logging.exception("Failed to abort prepared LeRobot files.") + if not close_started: + dataset.close() + raise + + +def _write_dataset_ray(table, dataset, info, source, source_schema, batch_size, + metadata, video_fields=(), *, parallelism, + source_options=None): + ray = _require_ray() + groups = _task_groups(info, dataset._episodes, video_fields, parallelism) + # Share only validated metadata, never the driver's dataset or open streams. + table_ref = ray.put(table.raw_table) + info_ref = ray.put(info) + metadata_ref = ray.put({ + "episodes": metadata["episodes"], + "subtask_indices": metadata["subtask_indices"], + }) + prepare = ray.remote(num_cpus=1, max_retries=0)(_prepare_group) + pending = [] + messages = [] + try: + for group in groups: + pending.append(prepare.remote( + table_ref, source.path, source_options, info_ref, + source_schema, batch_size, metadata_ref, video_fields, group)) + while pending: + ready, pending = ray.wait(pending, num_returns=1) + messages.extend(ray.get(ready[0])) Review Comment: [P2] Keep the ready task tracked until its messages are accepted `ray.wait` removes the ready task from `pending` before its messages are added to `messages`. If the driver receives Ctrl-C at this handoff, the exception handler only cancels/drains the remaining `pending` refs, so the successful task's prepared files are never aborted. With a real local Ray cluster, injecting `KeyboardInterrupt` during result delivery left three uncommitted Parquet, Blob, and Vector files; manually aborting the omitted messages removed them. This worker had already returned successfully, so this is separate from the documented lost-worker case. Please retain the current ready ref until its messages are accepted and include it in failure cleanup. ########## paimon-python/pypaimon/multimodal/lerobot/loader.py: ########## @@ -75,18 +77,50 @@ def _write_dataset( batch_size, metadata, video_fields=()): - target_schema = _target_schema(table.raw_table) - write_builder = table.raw_table.new_batch_write_builder() + messages = _prepare_dataset( + table.raw_table, dataset, info, source, source_schema, batch_size, + metadata, video_fields) + return _commit_messages(table.raw_table, messages) + + +def _commit_messages(table, messages): + snapshot_recorder = _SnapshotRecorder() + table_commit = table.new_batch_write_builder().new_commit() Review Comment: [P2] Initialize the committer before preparing data files `new_commit()` now runs after all frame files have been prepared, and its initialization failure leaves those messages unhandled. For example, `options={"commit.max-retries": "typo"}` is accepted when creating the frames table but fails when `FileStoreCommit` parses the option here. Using the actual writer, the base revision failed before writing frames and left zero files; this revision wrote the dataset and left three prepared files. Calling `_abort_messages` afterward cannot recover this configuration error because it constructs the same failing committer. Please initialize or preflight the driver committer before preparing files or scheduling workers, while preserving the existing handling of uncertain publication outcomes. ########## paimon-python/pypaimon/multimodal/lerobot/ray_loader.py: ########## @@ -0,0 +1,166 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ray workers prepare LeRobot files; only the driver commits snapshots.""" + +import logging +from pathlib import Path + +from pypaimon.multimodal.lerobot.loader import ( + _abort_messages, _commit_messages, _prepare_dataset, +) +from pypaimon.multimodal.lerobot.source import ( + _RemoteLeRobotDataset, _resolved_source, +) + + +_CANCEL_DRAIN_SECONDS = 1.0 + + +def _require_ray(): + try: + import ray + except ImportError as error: + raise ImportError("Parallel import requires 'pypaimon[ray,lerobot]'.") from error + if not ray.is_initialized(): + raise RuntimeError("Call ray.init() before using parallelism > 1.") + return ray + + +def _ray_source_uri(source): + if isinstance(source, str) and "://" in source: + return source + if isinstance(source, Path) or isinstance(source, str) and Path(source).is_dir(): + return Path(source).expanduser().resolve().as_uri() + raise ValueError( + "Parallel import requires a FileIO URI or shared directory, not a Hub repo_id.") + + +def _episode_groups(info, episodes, video_fields): + # Keep Episodes sharing a data or video file in the same task. + parents = list(range(len(episodes))) + + def root(i): + while parents[i] != i: + parents[i] = parents[parents[i]] + i = parents[i] + return i + + owners = {} + for i in range(len(episodes)): + episode = episodes[i] + data_path = _RemoteLeRobotDataset._data_path(episode, info) + previous = owners.setdefault(("data", data_path), i) + parents[root(i)] = root(previous) + for field in video_fields: + prefix = "videos/%s/" % field + path = info["video_path"].format( + video_key=field, chunk_index=episode[prefix + "chunk_index"], + file_index=episode[prefix + "file_index"]) + previous = owners.setdefault(("video", path), i) + parents[root(i)] = root(previous) + groups = {} + for i in range(len(episodes)): + groups.setdefault(root(i), []).append(i) + return list(groups.values()) + + +def _task_groups(info, episodes, video_fields, parallelism): + groups = _episode_groups(info, episodes, video_fields) + tasks = [[] for _ in range(min(parallelism, len(groups)))] + for i, group in enumerate(groups): + tasks[i % len(tasks)].extend(group) + return [sorted(task) for task in tasks] + + +def _prepare_group(table, source_path, source_options, info, schema, + batch_size, metadata, video_fields, episode_indices): + with _resolved_source(source_path, source_options) as (source, current_info): + if current_info != info: + raise ValueError("LeRobot source metadata changed during import.") + dataset = _RemoteLeRobotDataset(source, info) + messages = None + close_started = False + try: + messages = _prepare_dataset( + table, dataset, info, source, schema, batch_size, + metadata, video_fields, episode_indices) + close_started = True + dataset.close() + return messages Review Comment: [P2] Include source-context exit in prepared-file cleanup Returning here still executes the outer `_resolved_source` context's exit, which is outside the inner `try/except`. If soft cancellation raises `KeyboardInterrupt` while the source FileIO is closing, `_close_quietly` does not catch it, the worker does not return its messages, and this prepared-file abort handler is bypassed. I reproduced this with the real `_prepare_group` and writer by interrupting source-context exit after preparation: no abort ran and three uncommitted files remained. Please wrap the entire source context in the cleanup handler and retain ownership of the prepared messages until source closure has completed. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
