This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 9c7deebbde [python] Add distributed RoboMIND action backfill (#9496)
9c7deebbde is described below
commit 9c7deebbdeae9fde1945977163b6989e73172e77
Author: Yann Byron <[email protected]>
AuthorDate: Tue Sep 1 10:42:03 2026 +0800
[python] Add distributed RoboMIND action backfill (#9496)
---
docs/docs/pypaimon/robomind-agilex.md | 28 ++--
paimon-python/pypaimon/sample/robomind_agilex.py | 183 +++++++++++++++++----
.../tests/robomind_agilex_pipeline_test.py | 124 +++++++++++++-
3 files changed, 286 insertions(+), 49 deletions(-)
diff --git a/docs/docs/pypaimon/robomind-agilex.md
b/docs/docs/pypaimon/robomind-agilex.md
index e8aa3308d6..ee8aa65a10 100644
--- a/docs/docs/pypaimon/robomind-agilex.md
+++ b/docs/docs/pypaimon/robomind-agilex.md
@@ -84,9 +84,11 @@ pytest -q pypaimon/tests/robomind_agilex_pipeline_test.py \
```python
from pypaimon.sample.robomind_agilex import (
backfill_canonical_action,
+ backfill_canonical_action_ray,
ingest_local,
ingest_ray,
run_local_pipeline,
+ run_ray_pipeline,
)
# Run local ingestion and canonical-action backfill together.
@@ -95,25 +97,31 @@ pipeline = run_local_pipeline(
"/data/warehouse",
)
-# Or compose the lower-level operations explicitly. Ray chooses distributed
-# task placement; concurrency is only an optional upper bound.
-ingest = ingest_ray(
+# Run distributed ingestion and backfill on one managed Ray cluster.
+pipeline = run_ray_pipeline(
"/data/RoboMIND/h5_agilex_3rgb",
"/data/warehouse",
concurrency=8,
-)
-
-backfill = backfill_canonical_action(
- "/data/warehouse",
statistics_version="robomind-agilex-joint-position@1",
+ num_partitions=8,
+ ray_address="ray://cluster:10001",
)
```
Episode and frame ingestion commit separately and use the generic
`pypaimon.ray.load_from_hdf5` API in Ray mode. Canonical action materialization
-and statistics refresh also commit separately. If statistics need to be
-regenerated, call `refresh_action_statistics` without repeating ingestion or
-the row-id update.
+and statistics refresh also commit separately. The Ray backfill uses the
+optimized self-merge path: each target file group is processed by one task, so
+updates originating from multiple input batches cannot produce competing delta
+files for the same target file. The driver coordinates one commit after all
+file groups finish. If statistics need to be regenerated, call
+`refresh_action_statistics` without repeating ingestion or the row-id update.
+
+Use `backfill_canonical_action` for the local iterable path and
+`run_ray_pipeline` for managed distributed ingestion and backfill. The Ray
+pipeline requires Ray 2.50 or newer. The lower-level `ingest_ray` and
+`backfill_canonical_action_ray` stages assume that Ray has already been
+initialized, which allows either stage to be retried independently.
The canonical `action` is `float32(concat(master/joint_position_left,
master/joint_position_right))`. The backfill materializes only this consumed
diff --git a/paimon-python/pypaimon/sample/robomind_agilex.py
b/paimon-python/pypaimon/sample/robomind_agilex.py
index aa56133e25..5a4b83a170 100644
--- a/paimon-python/pypaimon/sample/robomind_agilex.py
+++ b/paimon-python/pypaimon/sample/robomind_agilex.py
@@ -114,6 +114,14 @@ class LocalPipelineResult:
backfill: BackfillResult
+@dataclass(frozen=True)
+class RayPipelineResult:
+ """Result of the complete Ray ingestion and backfill pipeline."""
+
+ ingest: IngestResult
+ backfill: BackfillResult
+
+
def episode_schema():
"""Return the shared AgileX episode business schema."""
return pa.schema([
@@ -341,25 +349,37 @@ def run_local_pipeline(
return LocalPipelineResult(ingest=ingest, backfill=backfill)
-def ingest_ray(
+def _require_ray_250(ray):
+ from packaging.version import parse
+
+ if parse(ray.__version__) < parse("2.50.0"):
+ raise RuntimeError(
+ "RoboMIND Ray backfill requires ray>=2.50; installed ray is %s."
+ % ray.__version__)
+
+
+def run_ray_pipeline(
input_root,
warehouse,
*,
database=DEFAULT_DATABASE,
batch_size=64,
concurrency=None,
+ statistics_version=DEFAULT_STATISTICS_VERSION,
+ num_partitions=None,
ray_address=None):
- """Ingest AgileX through the public distributed HDF5 loader."""
- if concurrency is not None:
- concurrency = _positive_int(concurrency, "concurrency")
- episodes = discover_episodes(input_root)
- _create_tables(warehouse, database)
+ """Run Ray ingestion and backfill on one managed Ray cluster."""
+ if num_partitions is not None:
+ num_partitions = _positive_int(num_partitions, "num_partitions")
+ if not isinstance(statistics_version, str) or not statistics_version:
+ raise ValueError("statistics_version must be a non-empty string.")
try:
import ray
except ImportError:
raise ImportError(
- "Ray ingestion requires ray; install pypaimon[ray,hdf5].")
+ "Ray pipeline requires ray; install pypaimon[ray,hdf5].")
+ _require_ray_250(ray)
initialized_here = not ray.is_initialized()
if initialized_here:
@@ -377,33 +397,62 @@ def ingest_ray(
"ray_address cannot be set after Ray has already been
initialized.")
try:
- from pypaimon.ray import load_from_hdf5
- catalog_options = {
- "warehouse": str(Path(warehouse).expanduser().resolve())}
- paths = [episode.path for episode in episodes]
- episode_result = load_from_hdf5(
- "%s.%s" % (database, EPISODES_TABLE), paths, catalog_options,
- transform=RoboMindAgileXEpisodeTransform(episodes),
+ ingest = ingest_ray(
+ input_root,
+ warehouse,
+ database=database,
+ batch_size=batch_size,
concurrency=concurrency,
)
- frame_result = load_from_hdf5(
- "%s.%s" % (database, FRAMES_TABLE), paths, catalog_options,
- transform=RoboMindAgileXFrameTransform(
- episodes, batch_size=batch_size),
- concurrency=concurrency,
- )
- return IngestResult(
- mode="ray",
- episode_count=episode_result.row_count,
- frame_count=frame_result.row_count,
- episodes_snapshot_id=episode_result.snapshot_id,
- frames_snapshot_id=frame_result.snapshot_id,
+ backfill = backfill_canonical_action_ray(
+ warehouse,
+ database=database,
+ statistics_version=statistics_version,
+ num_partitions=num_partitions,
)
+ return RayPipelineResult(ingest=ingest, backfill=backfill)
finally:
if initialized_here:
ray.shutdown()
+def ingest_ray(
+ input_root,
+ warehouse,
+ *,
+ database=DEFAULT_DATABASE,
+ batch_size=64,
+ concurrency=None):
+ """Ingest AgileX through an already initialized Ray cluster."""
+ if concurrency is not None:
+ concurrency = _positive_int(concurrency, "concurrency")
+ episodes = discover_episodes(input_root)
+ _create_tables(warehouse, database)
+
+ from pypaimon.ray import load_from_hdf5
+ catalog_options = {
+ "warehouse": str(Path(warehouse).expanduser().resolve())}
+ paths = [episode.path for episode in episodes]
+ episode_result = load_from_hdf5(
+ "%s.%s" % (database, EPISODES_TABLE), paths, catalog_options,
+ transform=RoboMindAgileXEpisodeTransform(episodes),
+ concurrency=concurrency,
+ )
+ frame_result = load_from_hdf5(
+ "%s.%s" % (database, FRAMES_TABLE), paths, catalog_options,
+ transform=RoboMindAgileXFrameTransform(
+ episodes, batch_size=batch_size),
+ concurrency=concurrency,
+ )
+ return IngestResult(
+ mode="ray",
+ episode_count=episode_result.row_count,
+ frame_count=frame_result.row_count,
+ episodes_snapshot_id=episode_result.snapshot_id,
+ frames_snapshot_id=frame_result.snapshot_id,
+ )
+
+
def build_canonical_action_backfill(source):
"""Build canonical actions keyed by the physical row ID."""
required = [_ACTION_LEFT, _ACTION_RIGHT, "_ROW_ID"]
@@ -453,8 +502,85 @@ def backfill_canonical_action(
)
+def backfill_canonical_action_ray(
+ warehouse,
+ *,
+ database=DEFAULT_DATABASE,
+ statistics_version=DEFAULT_STATISTICS_VERSION,
+ num_partitions=None):
+ """Run distributed backfill on an already initialized Ray cluster."""
+ try:
+ import ray
+ except ImportError:
+ raise ImportError(
+ "Ray backfill requires ray; install pypaimon[ray].")
+ _require_ray_250(ray)
+
+ row_count, frames_snapshot_id = _materialize_canonical_action_ray(
+ warehouse,
+ database=database,
+ num_partitions=num_partitions,
+ )
+ statistics_snapshot_id = refresh_action_statistics(
+ warehouse,
+ database=database,
+ statistics_version=statistics_version,
+ )
+ return BackfillResult(
+ row_count=row_count,
+ frames_snapshot_id=frames_snapshot_id,
+ statistics_snapshot_id=statistics_snapshot_id,
+ statistics_version=statistics_version,
+ )
+
+
def materialize_canonical_action(warehouse, *, database=DEFAULT_DATABASE):
"""Stage one: add and populate canonical action, then commit it."""
+ connection, frames_table = _prepare_canonical_action_table(
+ warehouse, database)
+ row_count = _update_canonical_action_batches(frames_table.raw_table)
+ frames_table = connection.get_table(FRAMES_TABLE)
+ frames_snapshot_id = _snapshot_id(frames_table)
+ return row_count, frames_snapshot_id
+
+
+def _materialize_canonical_action_ray(
+ warehouse,
+ *,
+ database=DEFAULT_DATABASE,
+ num_partitions=None):
+ """Stage one: materialize canonical action with Ray self-merge."""
+ if num_partitions is not None:
+ num_partitions = _positive_int(num_partitions, "num_partitions")
+
+ from pypaimon.ray import WhenMatched, merge_into
+
+ connection, _ = _prepare_canonical_action_table(
+ warehouse, database)
+ catalog_options = {
+ "warehouse": str(Path(warehouse).expanduser().resolve())}
+
+ def canonical_action(rows):
+ return build_canonical_action_backfill(rows)[_ACTION_COLUMN]
+
+ target = "%s.%s" % (database, FRAMES_TABLE)
+ result = merge_into(
+ target,
+ target,
+ catalog_options,
+ on=["_ROW_ID"],
+ when_matched=[WhenMatched.update({
+ _ACTION_COLUMN: canonical_action,
+ })],
+ read_columns=[_ACTION_LEFT, _ACTION_RIGHT],
+ num_partitions=num_partitions,
+ )
+
+ frames_table = connection.get_table(FRAMES_TABLE)
+ return result["num_matched"], _snapshot_id(frames_table)
+
+
+def _prepare_canonical_action_table(warehouse, database):
connection = pmm.connect(
database=database,
options={"warehouse": str(Path(warehouse).expanduser().resolve())},
@@ -478,10 +604,7 @@ def materialize_canonical_action(warehouse, *,
database=DEFAULT_DATABASE):
False,
)
frames_table = connection.get_table(FRAMES_TABLE)
- row_count = _update_canonical_action_batches(frames_table.raw_table)
- frames_table = connection.get_table(FRAMES_TABLE)
- frames_snapshot_id = _snapshot_id(frames_table)
- return row_count, frames_snapshot_id
+ return connection, frames_table
def refresh_action_statistics(
diff --git a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py
b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py
index 92013a9017..f77e146515 100644
--- a/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py
+++ b/paimon-python/pypaimon/tests/robomind_agilex_pipeline_test.py
@@ -15,6 +15,7 @@
# limitations under the License.
import importlib.util
+import inspect
import json
import subprocess
import sys
@@ -398,17 +399,18 @@ def
test_ray_ingest_matches_local_schema_rows_and_backfill(
local_backfill = agilex.backfill_canonical_action(
local_warehouse, statistics_version="synthetic-actions@1")
- ray.init(num_cpus=2, include_dashboard=False)
- try:
- ray_result = agilex.ingest_ray(
- root, ray_warehouse, batch_size=2, concurrency=2)
- finally:
- ray.shutdown()
+ ray_pipeline = agilex.run_ray_pipeline(
+ root,
+ ray_warehouse,
+ batch_size=2,
+ concurrency=2,
+ statistics_version="synthetic-actions@1",
+ )
+ ray_result = ray_pipeline.ingest
+ ray_backfill = ray_pipeline.backfill
+ assert not ray.is_initialized()
assert ray_result.episodes_snapshot_id == 1
assert ray_result.frames_snapshot_id == 1
- ray_backfill = agilex.backfill_canonical_action(
- ray_warehouse, statistics_version="synthetic-actions@1")
-
assert ray_result.mode == "ray"
assert ray_result.episode_count == len(paths)
assert ray_result.frame_count == 12
@@ -426,6 +428,110 @@ def
test_ray_ingest_matches_local_schema_rows_and_backfill(
_logical_rows(ray_warehouse, table_name, schema))
+def test_run_ray_pipeline_owns_cluster_lifecycle(monkeypatch):
+ ray = pytest.importorskip("ray")
+ ray_init = MagicMock()
+ ray_shutdown = MagicMock()
+ ingest = MagicMock(return_value=MagicMock())
+ backfill = MagicMock(return_value=MagicMock())
+ monkeypatch.setattr(ray, "__version__", "2.50.0")
+ monkeypatch.setattr(ray, "is_initialized", lambda: False)
+ monkeypatch.setattr(ray, "init", ray_init)
+ monkeypatch.setattr(ray, "shutdown", ray_shutdown)
+ monkeypatch.setattr(agilex, "ingest_ray", ingest)
+ monkeypatch.setattr(agilex, "backfill_canonical_action_ray", backfill)
+
+ result = agilex.run_ray_pipeline(
+ "input", "warehouse", ray_address="ray://cluster")
+
+ ray_init.assert_called_once_with(
+ include_dashboard=False,
+ ignore_reinit_error=True,
+ address="ray://cluster",
+ )
+ ingest.assert_called_once_with(
+ "input",
+ "warehouse",
+ database=agilex.DEFAULT_DATABASE,
+ batch_size=64,
+ concurrency=None,
+ )
+ backfill.assert_called_once_with(
+ "warehouse",
+ database=agilex.DEFAULT_DATABASE,
+ statistics_version=agilex.DEFAULT_STATISTICS_VERSION,
+ num_partitions=None,
+ )
+ ray_shutdown.assert_called_once_with()
+ assert result.ingest is ingest.return_value
+ assert result.backfill is backfill.return_value
+
+
+def test_run_ray_pipeline_rejects_old_ray_before_ingest(monkeypatch):
+ ray = pytest.importorskip("ray")
+ ingest = MagicMock()
+ monkeypatch.setattr(ray, "__version__", "2.49.0")
+ monkeypatch.setattr(ray, "is_initialized", lambda: True)
+ monkeypatch.setattr(agilex, "ingest_ray", ingest)
+
+ with pytest.raises(RuntimeError, match="requires ray>=2.50"):
+ agilex.run_ray_pipeline("input", "warehouse")
+
+ ingest.assert_not_called()
+
+
+def test_ray_backfill_rejects_old_ray_before_preparing_table(monkeypatch):
+ ray = pytest.importorskip("ray")
+ prepare = MagicMock(return_value=(MagicMock(), MagicMock()))
+ monkeypatch.setattr(ray, "__version__", "2.49.0")
+ monkeypatch.setattr(agilex, "_prepare_canonical_action_table", prepare)
+
+ with pytest.raises(RuntimeError, match="requires ray>=2.50"):
+ agilex.backfill_canonical_action_ray("warehouse")
+
+ prepare.assert_not_called()
+
+
+def test_run_ray_pipeline_rejects_num_partitions_before_ingest(monkeypatch):
+ ray = pytest.importorskip("ray")
+ ingest = MagicMock()
+ monkeypatch.setattr(ray, "__version__", "2.50.0")
+ monkeypatch.setattr(ray, "is_initialized", lambda: True)
+ monkeypatch.setattr(agilex, "ingest_ray", ingest)
+
+ with pytest.raises(
+ ValueError, match="num_partitions must be a positive int"):
+ agilex.run_ray_pipeline(
+ "input", "warehouse", num_partitions=0)
+
+ ingest.assert_not_called()
+
+
+def test_run_ray_pipeline_rejects_statistics_version_before_ingest(
+ monkeypatch):
+ ray = pytest.importorskip("ray")
+ ingest = MagicMock()
+ monkeypatch.setattr(ray, "__version__", "2.50.0")
+ monkeypatch.setattr(ray, "is_initialized", lambda: True)
+ monkeypatch.setattr(agilex, "ingest_ray", ingest)
+ monkeypatch.setattr(
+ agilex, "_materialize_canonical_action_ray",
+ MagicMock(return_value=(12, 2)))
+
+ with pytest.raises(
+ ValueError, match="statistics_version must be a non-empty string"):
+ agilex.run_ray_pipeline(
+ "input", "warehouse", statistics_version="")
+
+ ingest.assert_not_called()
+
+
+def test_ray_stage_apis_do_not_accept_cluster_addresses():
+ assert "ray_address" not in inspect.signature(agilex.ingest_ray).parameters
+ assert "ray_address" not in inspect.signature(
+ agilex.backfill_canonical_action_ray).parameters
+
+
def
test_stream_action_statistics_are_stable_and_order_independent(monkeypatch):
frame_count = 257
row = np.arange(frame_count, dtype=np.float64).reshape(-1, 1)
