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_dataset_v21_to_v30_udf import ConvertDatasetV21ToV30
9
10if __name__ == "__main__":
11    if os.getenv("DAFT_RUNNER", "native") == "ray":
12        import ray
13        ray.init(dashboard_host="0.0.0.0", ignore_reinit_error=True)
14        daft.set_runner_ray()
15    daft.set_execution_config(actor_udf_ready_timeout=6000, min_cpu_per_task=0)
16
17    tasks = {
18        "input_repoid": [
19            "lerobot/aloha_sim_insertion_human/", 
20            "lerobot/pusht/",
21            "lerobot/pusht2/"   
22        ],
23        "input_path": ["/mnt/pfs/xxx/dataset/lerobot_v2/"] * 3,
24        "output_path": ["/mnt/pfs/xxx/dataset/lerobot_v3_output/"] * 3
25    }
26    num_datasets = len(tasks)
27 
28    ds = daft.from_pydict(tasks)
29    ds = ds.into_partitions(num_datasets)
30
31    ds = ds.with_column(
32        "convert_result",
33        aihc_udf(
34            ConvertDatasetV21ToV30,
35            construct_args={
36            },
37            num_cpus=1,
38            batch_size=1,
39            concurrency=num_datasets,
40            use_process=True
41        )(col("input_repoid"), col("input_path"), col("output_path")),
42    )
43
44    ds.show()

1#启动命令
2#v2tov3_udf.py 为以上的示例代码
3
4#!/bin/bash
5echo "Start Execution."
6python /mnt/pfs/cxy/test/aihc/embodied_ai/v2_to_v3/v2tov3_opreator.py
7echo "Execution completed."