1from __future__ import annotations
2
3import os
4import daft
5from daft import col
6
7from daft.aihc.common.udf import aihc_udf
8from daft.aihc.functions.embodied.convert_rlds_to_lerobot_udf import ConvertRldsToLerobot
9
10os.environ["DAFT_RUNNER"] = "ray"
11
12if __name__ == "__main__":
13    if os.getenv("DAFT_RUNNER", "native") == "ray":
14        import logging
15        import ray
16
17        try:
18            ray.init(address="auto", ignore_reinit_error=True)
19        except ConnectionError:
20            ray.init(ignore_reinit_error=True)
21
22    daft.set_execution_config(
23        actor_udf_ready_timeout=6000,
24        min_cpu_per_task=0,
25    )
26
27    # === 配置 ===
28    DATASET_NAME = "cmu_play_fusion"
29    DATA_DIR = "/mnt/pfs/data/oxe/"
30    OUTPUT_PATH = "/mnt/pfs/data/v3_output/cmu_play_fusion_lerobot"
31    NUM_SPLITS = 5
32
33    # 输入: split_idx 列表
34    tasks = {
35        "split_idx": list(range(NUM_SPLITS)),
36    }
37
38    ds = daft.from_pydict(tasks)
39    ds = ds.into_partitions(NUM_SPLITS)
40
41    ds = ds.with_column(
42        "convert_result",
43        aihc_udf(
44            ConvertRldsToLerobot,
45            construct_args={
46                "dataset_name": DATASET_NAME,
47                "data_dir": DATA_DIR,
48                "output_path": OUTPUT_PATH,
49                "num_splits": NUM_SPLITS,
50                "use_videos": True,
51                "image_writer_processes": 5,
52                "image_writer_threads": 10,
53            },
54            num_cpus=1,
55            batch_size=1,
56            concurrency=NUM_SPLITS,
57            use_process=True,
58        )(col("split_idx")),
59    )
60
61    ds.show()
62
63    # === 合并所有 split 子集 ===
64    if NUM_SPLITS > 1:
65        from daft.aihc.functions.embodied.convert_rlds_to_lerobot_udf import merge_splits
66        merge_splits(
67            output_path=OUTPUT_PATH,
68            repo_id=DATASET_NAME,
69            num_splits=NUM_SPLITS,
70            cleanup=True,
71        )