rbs-data-utils 0.1.0
Foundational utility library for RBS data pipelines.
Requires Python: >=3.10

# rbs-data-utils

`rbs-data-utils` is the foundational utility library for RBS data pipelines.

It provides:
- Task-scoped artifact isolation
- Concurrency-safe JSONL publishing (part -> publish jsonl)
- Standardized directory layout contracts
- Incremental publish readers
- Sharding and worker utilities

This repository does NOT contain model logic.
It defines the infrastructure layer that guarantees data consistency,
traceability, and scalability across pipeline modules.

## Package layout

```text
rbs-data-utils/
├── README.md
├── pyproject.toml
├── src/
│   ├── __init__.py
│   ├── io/
│   │   ├── io_manager.py
│   │   ├── errors.py
│   │   └── file_lock.py
│   ├── layout/
│   ├── schema/
│   ├── index/
│   ├── concurrency/
│   └── utils/
├── demo/
└── docs/
```

> Note: 代码实现放在 `src/`，安装后外部通过 `import rbs_data_utils` 使用 API。

## Install and import

```bash
pip install -e .
```

```python
from rbs_data_utils import IOManager

# 1) 初始化任务与模块。产物从 /tmp/io_manager_tasks/<task_name>/... 读写
dataset_name = "demo_dataset_a"
target_task_name = "demo_task"
mgr = IOManager(
    task_name=target_task_name,
    module_names=["vlm_process"],
    base_dir="/tmp/io_manager_tasks",
)

# 2) 不传 worker_key，IOManager 会默认生成进程级 worker key
# 写入两条样例记录到当前 worker 的 part 文件
mgr.write_part(
    module="vlm_process",
    record={"id": 1, "status": "ok"},
    dataset_name=dataset_name,
    task_name=target_task_name,
)
mgr.write_part(
    module="vlm_process",
    record={"id": 2, "status": "ok"},
    dataset_name=dataset_name,
    task_name=target_task_name,
)

# 3) 合并当前 worker 的 part 到 <dataset>_publish.jsonl（默认 remove_part=True）
merge_result = mgr.write_finished(
    module="vlm_process",
    dataset_name=dataset_name,
    task_name=target_task_name,
)

# 4) 读取 publish：全量、偏移、偏移+限制
publish_all = mgr.read_publish(
    module="vlm_process",
    dataset_name=dataset_name,
    task_name=target_task_name,
)
publish_offset_1 = mgr.read_publish(
    module="vlm_process",
    dataset_name=dataset_name,
    task_name=target_task_name,
    offset=1,
)
publish_offset_1_limit_1 = mgr.read_publish(
    module="vlm_process",
    dataset_name=dataset_name,
    task_name=target_task_name,
    offset=1,
    limit=1,
)

# 5) 结果变量说明
# - merge_result: 合并统计信息（merged/skipped/throttled）
# - publish_all: 全量读取结果
# - publish_offset_1: 从第 2 条开始读取
# - publish_offset_1_limit_1: 从第 2 条开始只读取 1 条
```

## Best practice: full pipeline task (`demo_all`)

下面以一次全量生产任务 `demo_all` 为例，模块包含：
`vlm_anno`、`vlm_seg`、`sam3`、`sam3d`、`any6d`、`any4d`、`fp++`。

假设数据根目录（`IOManager.base_dir`）固定为：

```text
/demo_dataset/
```

建议每个模块都用**同一组 task 配置**初始化自己的 `IOManager`，并且只声明本模块名；这样更贴近日常生产写法，同时仍然落在同一个 task 空间：

```python
from rbs_data_utils import IOManager

TASK_NAME = "demo_all"
BASE_DIR = "/demo_dataset"


def make_mgr(self_module: str) -> IOManager:
    # 每个模块进程独立初始化，只注册自己模块
    # IOManager 会自动创建并复用目录（存在则不重复创建）
    return IOManager(
        task_name=TASK_NAME,
        module_names=[self_module],
        base_dir=BASE_DIR,
    )
```

### Recommended module data flow

```text
vlm_anno -> vlm_seg -> sam3 -> sam3d -> any6d -> any4d -> fp++
```

- 读上游：`mgr.read_publish(module="<upstream_module>", dataset_name=dataset_name, task_name=upstream_task_name)`
- 写本模块增量结果：`mgr.write_part(module="<self_module>", record=record, dataset_name=dataset_name, task_name=self_task_name)`
- 完成一批后合并发布：`mgr.write_finished(module="<self_module>", dataset_name=dataset_name, task_name=self_task_name)`（默认删除已合并 part）

示例（以 `sam3d` 模块为例）：

```python
def run_sam3d():
    dataset_name = "AgiBotWorld-Beta"
    upstream_task_name = "demo_all_v1"
    self_task_name = "demo_all_v2"
    mgr = make_mgr(self_module="sam3d")

    # 1) 读取上游 sam3 的 <dataset>_publish.jsonl
    sam3_records = mgr.read_publish(
        module="sam3",
        dataset_name=dataset_name,
        task_name=upstream_task_name,
    )

    # 2) 逐条处理后写入 sam3d 的 part 文件
    for rec in sam3_records:
        out = {
            "sample_id": rec.get("sample_id"),
            "sam3d_result": "ok",
        }
        mgr.write_part(
            module="sam3d",
            record=out,
            dataset_name=dataset_name,
            task_name=self_task_name,
        )

    # 3) 将本 worker 的 part 合并到 sam3d/<dataset>_publish.jsonl
    merge_info = mgr.write_finished(
        module="sam3d",
        dataset_name=dataset_name,
        task_name=self_task_name,
        remove_part=True,
    )
    print("sam3d merge:", merge_info)
```

### Upstream / self-module JSONL usage template

你可以用同一模板复用到每个模块（把 `SELF` 和 `UPSTREAMS` 改成对应模块）：

```python
SELF = "any6d"
DATASET_NAME = "AgiBotWorld-Beta"
UPSTREAM_TASK_NAME = "demo_all_v1"
SELF_TASK_NAME = "demo_all_v2"
UPSTREAMS = ["sam3d"]  # 多上游时可写成 ["sam3d", "vlm_seg"]

mgr = make_mgr(self_module=SELF)

# 读上游 publish.jsonl
inputs = []
for up in UPSTREAMS:
    inputs.extend(
        mgr.read_publish(module=up, dataset_name=DATASET_NAME, task_name=UPSTREAM_TASK_NAME)
    )

# 写本模块 part.jsonl
for item in inputs:
    result = {"sample_id": item.get("sample_id"), "module": SELF}
    mgr.write_part(
        module=SELF,
        record=result,
        dataset_name=DATASET_NAME,
        task_name=SELF_TASK_NAME,
    )

# 发布本模块 publish.jsonl
mgr.write_finished(module=SELF, dataset_name=DATASET_NAME, task_name=SELF_TASK_NAME)

# 可选：回读本模块 publish，做抽样检查
print(
    mgr.read_publish(
        module=SELF,
        dataset_name=DATASET_NAME,
        task_name=SELF_TASK_NAME,
        limit=5,
    )
)
```

### Practical suggestions

- 默认按“各模块仅注册自己”即可；如需固定 manifest，也可在任务启动时一次性创建全量 `module_names`。
- 不需要手动传 `worker_key`，默认即可（底层会生成进程 worker key）。
- `write_finished` 可按批次调用，不必等全量结束；库会自动做并发安全合并与去重。
- 需要存放中间产物（图片、npy、可视化）时，使用 `artifact_path(module, *subpaths)` 获取标准路径。

## Run demos

```bash
pip install -e .
python demo/demo.py
```