perf(datasets): cache OSS bucket, add server-side pagination and --filter for tasks

[xdlkc]

fixes #1063

Summary

• Cache oss2.Bucket instance — avoids creating 25+ HTTP sessions per list_all_datasets call (datasets list: 5.0s → 2.7s)
• Server-side pagination — push offset/limit down to _extract_tasks_from_split with early termination and continuation_token caching (datasets tasks --limit 10: 19.4s → 1.5s)
• --filter flag — prefix-based task search pushed down to OSS query prefix (datasets tasks --filter 0xerr0r: 1.4s)

Design

Bucket caching

_build_bucket() now lazily initializes and caches the oss2.Bucket instance. All methods on the same OssDatasetRegistry share one HTTP session.

Pagination cache

A _PaginationCache dataclass stores the last query's (query_prefix, tasks, continuation_token, is_exhausted). When the next query matches (same prefix, sequential page access), it resumes from the cached token instead of re-scanning from page 1. This turns sequential pagination (page 1 → 2 → 3) into O(1) per page.

max_keys is adapted per request: min(1000, items_still_needed) instead of always 1000.

Filter push-down

--filter is appended to the OSS query prefix (split_prefix + filter), so OSS only returns matching keys. Task names are extracted relative to the original split_prefix to preserve full names.

Interface changes

list_dataset_tasks gains keyword-only offset, limit, task_filter parameters across the full chain: BaseDatasetRegistry → OssDatasetRegistry → DatasetClient → CLI. All default to their zero-value, so existing callers are unaffected.

Benchmark

34,732 tasks, oss-ap-southeast-1, ~0.5s RTT.

Scenario	Before	After	Speedup
datasets list	5.0s	2.7s	1.9×
datasets tasks --limit 10	19.4s	1.5s	13×
datasets tasks (no limit)	19.4s	15.5s	1.3×
datasets tasks --filter prefix	N/A	1.4s	new

Test plan

• pytest tests/unit/datasets/ — 99 tests pass
• Manual verification: datasets list, datasets tasks --limit, datasets tasks --filter against rock-agent-pre bucket

Co-Authored-By: Claude Code noreply@anthropic.com

[alibabarock]

Code Review Summary

Verdict: Comment（无 blocking 问题，有几点建议）

⚠️ Warnings

1. _fs_download 目录下载时 get_task_file 被调了两次 — 先调一次探测"单文件还是目录"，失败后再 list + 逐个 get。对大目录来说浪费了一次 RTT。建议先 list 再判断，避免无意义的 get。

2. _PaginationCache 的 tasks 字段会缓存全量已扫描的 task_ids — 34k tasks 场景下内存里存了 34k 字符串。如果反复切不同 dataset/split 查询，前一个的 cache 会一直占着。建议加个 LRU 或至少 cache 过期。

3. --ouput typo 作为 alias 被保留 — add_argument("-o", "--output", "--ouput", ...)。理解是兼容拼写错误，但不太常规，建议注明或直接去掉。

💡 Suggestions

1. _extract_tasks_from_split 的 max_keys 自适应 — OSS 返回的 key 可能有重复（task 目录本身 + 子文件），去重后可能不够。建议 min(1000, items_still_needed + 10) 作为 buffer。

2. _upload_task_file 返回 int | None — None 表示 skipped，1 表示 uploaded，和 _upload_task 返回的 int 混在一起语义模糊。建议统一用 enum 或 dataclass 表示状态。

3. _normalize_task_path 路径安全校验 — 防 .. 做得好。. 和空段被过滤是正确行为，建议在 docstring 里说明。

✅ Looks Good

• Bucket 缓存 + 分页下推 + filter prefix 下推，性能提升显著（13x）
• _PaginationCache 设计合理，顺序翻页 O(1)
• datasets fs 子命令结构清晰，路径安全校验到位
• 隐藏文件过滤覆盖了 org/dataset/split 三层
• JSON 输出统一用 _print_json，contextlib.redirect_stdout 隔离 upload 噪音
• 99 个单测全过，覆盖充分

[StephenRi]

函数都封装到类里面