Skip to content

Bug fix when set total_steps. #386

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pan-x-c merged 4 commits into from
Nov 14, 2025
Merged

Bug fix when set total_steps. #386

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
c562a7e
Bug fix when set `total_steps`
chenyushuo Nov 14, 2025
fedc49a
apply suggestions
chenyushuo Nov 14, 2025
4462fed
refactor _should_stop
chenyushuo Nov 14, 2025
e717462
add unittest point
chenyushuo Nov 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 58 additions & 3 deletions tests/buffer/task_scheduler_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,54 @@ def _check_batch_tasks(self, batch_tasks: List[Task], indices: List[Dict[str, in
@parameterized.expand(
[
(
{"batch_size": 5, "total_steps": 3},
{"selector_type": "sequential"},
[
{"index": 0, "taskset_id": 1},
{"index": 1, "taskset_id": 1},
{"index": 2, "taskset_id": 1},
{"index": 0, "taskset_id": 0},
{"index": 1, "taskset_id": 0},
{"index": 3, "taskset_id": 1},
{"index": 4, "taskset_id": 1},
{"index": 5, "taskset_id": 1},
{"index": 2, "taskset_id": 0},
{"index": 3, "taskset_id": 0},
{"index": 6, "taskset_id": 1},
{"index": 0, "taskset_id": 1},
{"index": 1, "taskset_id": 1},
{"index": 4, "taskset_id": 0},
{"index": 0, "taskset_id": 0},
],
),
(
{"batch_size": 5, "total_epochs": 2},
{"selector_type": "sequential"},
[
{"index": 0, "taskset_id": 1},
{"index": 1, "taskset_id": 1},
{"index": 2, "taskset_id": 1},
{"index": 0, "taskset_id": 0},
{"index": 1, "taskset_id": 0},
{"index": 3, "taskset_id": 1},
{"index": 4, "taskset_id": 1},
{"index": 5, "taskset_id": 1},
{"index": 2, "taskset_id": 0},
{"index": 3, "taskset_id": 0},
{"index": 6, "taskset_id": 1},
{"index": 0, "taskset_id": 1},
{"index": 1, "taskset_id": 1},
{"index": 4, "taskset_id": 0},
{"index": 0, "taskset_id": 0},
{"index": 2, "taskset_id": 1},
{"index": 3, "taskset_id": 1},
{"index": 4, "taskset_id": 1},
{"index": 1, "taskset_id": 0},
{"index": 2, "taskset_id": 0},
],
),
(
{"batch_size": 2, "total_epochs": 2},
{"selector_type": "sequential"},
[
{"index": 0, "taskset_id": 1},
Expand Down Expand Up @@ -70,6 +118,7 @@ def _check_batch_tasks(self, batch_tasks: List[Task], indices: List[Dict[str, in
],
),
(
{"batch_size": 2, "total_epochs": 2},
{"selector_type": "shuffle", "seed": 42},
[
{"index": 3, "taskset_id": 1},
Expand Down Expand Up @@ -99,6 +148,7 @@ def _check_batch_tasks(self, batch_tasks: List[Task], indices: List[Dict[str, in
],
),
(
{"batch_size": 2, "total_epochs": 2},
{"selector_type": "random", "seed": 42},
[
{"index": 0, "taskset_id": 1},
Expand Down Expand Up @@ -128,6 +178,7 @@ def _check_batch_tasks(self, batch_tasks: List[Task], indices: List[Dict[str, in
],
),
(
{"batch_size": 2, "total_epochs": 2},
{"selector_type": "offline_easy2hard", "feature_keys": ["feature_offline"]},
[
{"index": 3, "taskset_id": 1},
Expand Down Expand Up @@ -157,6 +208,7 @@ def _check_batch_tasks(self, batch_tasks: List[Task], indices: List[Dict[str, in
],
),
(
{"batch_size": 2, "total_epochs": 2},
{"selector_type": "difficulty_based", "feature_keys": ["feat_1", "feat_2"]},
[
{"index": 3, "taskset_id": 1},
Expand Down Expand Up @@ -187,10 +239,13 @@ def _check_batch_tasks(self, batch_tasks: List[Task], indices: List[Dict[str, in
),
]
)
async def test_task_scheduler(self, task_selector_kwargs, batch_tasks_orders) -> None:
async def test_task_scheduler(
self, buffer_config_kwargs, task_selector_kwargs, batch_tasks_orders
) -> None:
config = get_template_config()
config.buffer.batch_size = 2
config.buffer.total_epochs = 2
config.mode = "explore"
for key, value in buffer_config_kwargs.items():
setattr(config.buffer, key, value)
config.buffer.explorer_input.taskset = None
config.buffer.explorer_input.tasksets = [
TasksetConfig(
Expand Down
20 changes: 12 additions & 8 deletions trinity/buffer/task_scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,13 @@ def __init__(self, explorer_state: Dict, config: Config):
self.epoch = self.step * self.read_batch_size // len(self.base_taskset_ids)
self.orders = self.build_orders(self.epoch)

if self.config.buffer.total_steps:
self.max_steps = self.config.buffer.total_steps
else:
self.max_steps = (
self.config.buffer.total_epochs * len(self.base_taskset_ids) // self.read_batch_size
)

def build_orders(self, epoch: int):
"""
Creates a shuffled sequence of taskset IDs to control sampling priority per step.
Expand All @@ -108,6 +115,9 @@ def build_orders(self, epoch: int):
rng.shuffle(taskset_ids)
return taskset_ids

def _should_stop(self) -> bool:
return self.step >= self.max_steps

async def read_async(self) -> List:
"""
Asynchronously reads a batch of tasks according to the current schedule.
Expand All @@ -125,12 +135,8 @@ async def read_async(self) -> List:
Returns:
List[Task]: A batch of tasks from potentially multiple tasksets
"""
if self.config.buffer.total_steps:
if self.step >= self.config.buffer.total_steps:
raise StopAsyncIteration
else:
if self.epoch >= self.config.buffer.total_epochs:
raise StopAsyncIteration
if self._should_stop():
raise StopAsyncIteration

batch_size = self.read_batch_size
start = self.step * batch_size % len(self.base_taskset_ids)
Expand All @@ -143,8 +149,6 @@ async def read_async(self) -> List:
else:
taskset_ids = self.orders[start:]
self.epoch += 1
if self.epoch >= self.config.buffer.total_epochs:
raise StopAsyncIteration
self.orders = self.build_orders(self.epoch)
taskset_ids += self.orders[: (end - len(self.base_taskset_ids))]

Expand Down