alibaba / alibaba/ROCK

[Feature] datasets tasks: server-side pagination, continuation token caching, and --filter support

Open
#1,063 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
485
Forks
81
Avg merge
16h 12m
Merged PRs (30d)
8

Description

**Feature Category**
- [ ] Sandbox
- [ ] Actions
- [ ] Deployments
- [x] SDK & API
- [ ] Envhub
- [x] CLI
- [x] Performance & Optimization
- [ ] Documentation & Examples

**Problem Statement**

The `datasets tasks` command has two performance issues and a missing search capability:

1. **Redundant Bucket creation**: `OssDatasetRegistry._build_bucket()` creates a new `oss2.Bucket` (and underlying HTTP session) on every method call. `list_all_datasets` creates 25+ Bucket instances — one for `list_organizations` and one per org for `list_org_datasets`. Measured impact: `datasets list` takes **5.0s** instead of **2.7s**.

2. **No server-side pagination**: `list_dataset_tasks` always scans all OSS pages serially (35 pages × ~0.5s RTT for 34,732 tasks = **19.4s**), then applies `--offset`/`--limit` as a client-side slice. Even `--limit 10` pays the full 19.4s scan cost.

3. **No task search**: Users cannot filter tasks by prefix. To find tasks for a specific repo (e.g. `0xerr0r__blocky-*`) they must list all 34,732 tasks and grep locally.

**Proposed Solution**

1. **Cache Bucket instance**: Lazy-initialize `oss2.Bucket` once and reuse across all method calls on the same `OssDatasetRegistry` instance.

2. **Push offset/limit to OSS layer**:
- Add `offset`, `limit`, `task_filter` keyword arguments to `list_dataset_tasks` (through `BaseDatasetRegistry` → `DatasetClient` → CLI).
- Compute `max_items = offset + limit` and pass to `_extract_tasks_from_split`, which stops paginating as soon as enough items are collected.
- Adapt `max_keys` per request to match the remaining item count instead of always requesting 1000.
- Cache the last query's `continuation_token` and accumulated tasks in a `_PaginationCache` dataclass. Sequential page access (page 1 → page 2 → page 3) resumes from the cached token in O(1) instead of re-scanning from the beginning.

3. **Add `--filter` to `datasets tasks`**:
- Push the filter string down as an OSS prefix (`split_prefix + filter`), so OSS only returns matching keys server-side.
- Extract task names relative to the original `split_prefix` to preserve full task names in output.

**Benchmark**

Environment: `oss-ap-southeast-1`, dataset with 34,732 tasks, ~0.5s RTT per request.

| Scenario | Before | After |
|---|---|---|
| `datasets list` | 5.0s | **2.7s** |
| `datasets tasks --limit 10` | 19.4s | **1.5s** |
| `datasets tasks` (full scan) | 19.4s | **15.5s** |
| `datasets tasks --filter 0xerr0r` | N/A | **1.4s** |

**Detailed Changes**

- `OssDatasetRegistry.__init__`: add `_bucket` cache and `_PaginationCache` instance
- `_build_bucket()`: lazy-init, return cached `oss2.Bucket`
- `_extract_tasks_from_split()`: accept `max_items` and `task_filter`; implement pagination cache with continuation token; adapt `max_keys` dynamically
- `list_dataset_tasks()`: accept and forward `offset`, `limit`, `task_filter`; compute `max_items` and slice result
- `BaseDatasetRegistry.list_dataset_tasks()`: add `offset`, `limit`, `task_filter` to abstract signature
- `DatasetClient.list_dataset_tasks()`: forward new parameters
- CLI `_tasks()`: pass `offset`/`limit`/`filter` to client instead of client-side slicing
- CLI parser: add `--filter` argument

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with OssDatasetRegistry.__init__, _build_bucket(), and _extract_tasks_from_split() to understand bucket reuse and pagination behavior. Then trace list_dataset_tasks() through BaseDatasetRegistry, DatasetClient, and the CLI _tasks() parser. Done means server-side offset, limit, and prefix filtering work, sequential pages reuse cached state, and the stated benchmark improvements are reproducible.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
api, cli, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.