trinity.explorer.scheduler module

Scheduler for rollout tasks.

class trinity.explorer.scheduler.TaskWrapper(task: ~trinity.common.workflows.workflow.Task, batch_id: int | str, sub_task_num: int = 1, finished_sub_task_num: int = 0, completed_runs: int = 0, total_runs: int = 0, metrics: ~typing.List[~typing.Dict[str, float]] = <factory>, experience_payloads: ~typing.List[bytes] = <factory>, first_error: str | None = None, emitted: bool = False)

Bases: object

A wrapper for a task. Each task can run multiple times (repeat_times) on same or different runners.

task: Task

batch_id: int | str

sub_task_num: int = 1

finished_sub_task_num: int = 0

completed_runs: int = 0

total_runs: int = 0

metrics: List[Dict[str, float]]

experience_payloads: List[bytes]

first_error: str | None = None

emitted: bool = False

__init__(task: ~trinity.common.workflows.workflow.Task, batch_id: int | str, sub_task_num: int = 1, finished_sub_task_num: int = 0, completed_runs: int = 0, total_runs: int = 0, metrics: ~typing.List[~typing.Dict[str, float]] = <factory>, experience_payloads: ~typing.List[bytes] = <factory>, first_error: str | None = None, emitted: bool = False) → None

class trinity.explorer.scheduler.CompletedTaskResult(batch_id: int | str, task_id: int | str, status: ~trinity.explorer.workflow_runner.Status, experience_payloads: ~typing.List[bytes] = <factory>)

Bases: object

A completed task result stored by batch and task id.

batch_id: int | str

task_id: int | str

status: Status

experience_payloads: List[bytes]

__init__(batch_id: int | str, task_id: int | str, status: ~trinity.explorer.workflow_runner.Status, experience_payloads: ~typing.List[bytes] = <factory>) → None

class trinity.explorer.scheduler.RunningTaskState(task: TaskWrapper, runner_id: int, restart_runner_on_cancel: bool = True)

Bases: object

Per-future execution state tracked while one task is running.

task: TaskWrapper

runner_id: int

restart_runner_on_cancel: bool = True

__init__(task: TaskWrapper, runner_id: int, restart_runner_on_cancel: bool = True) → None

class trinity.explorer.scheduler.RunnerWrapper(runner_id: int, rollout_model_id: int, auxiliary_model_ids: List[int], config: Config)

Bases: object

A wrapper for a WorkflowRunner

__init__(runner_id: int, rollout_model_id: int, auxiliary_model_ids: List[int], config: Config)

async update_state() → None

Get the runner state.

async prepare()

async run_with_retry(task: TaskWrapper, repeat_times: int, run_id_base: int, timeout: float, collect_partial_runs: bool) → Tuple[Status, bytes, int, float]

Parameters:

• task (TaskWrapper) – The task to run.

• repeat_times (int) – The number of times to repeat the task.

• run_id_base (int) – The base run id for this task runs.

• timeout (float) – The timeout for each task run.

Returns:

The return status of the task. List: The experiences generated by the task. int: The runner_id of current runner. float: The time taken to run the task.

Return type:

Status

async restart_runner()

trinity.explorer.scheduler.sort_batch_id(batch_id: int | str)

Priority of batch_id

class trinity.explorer.scheduler.Scheduler(config: Config)

Bases: object

Scheduler for rollout tasks.

Supports scheduling tasks to multiple runners, retrying failed tasks, and collecting results at different levels.

__init__(config: Config)

task_done_callback(async_task: Task)

discard_completed_results(batch_id: int | str) → None

Drop cached completed results for one batch.

async start() → None

async stop() → None

schedule(tasks: List[Task], batch_id: int | str) → None

Schedule the provided tasks.

Parameters:

• tasks (List[Task]) – The tasks to schedule.

• batch_id (Union[int, str]) – The id of provided tasks. In most cases, it should be current step number for training tasks and “<current_step_num>/<eval_taskset_name>” for eval tasks.

dynamic_timeout(timeout: float | None = None) → float

Calculate dynamic timeout based on historical data.

async drain_batch_payload_results(batch_id: int | str) → Tuple[List[Status], List[bytes]]

Drain cached completed results for one batch.

async get_payload_results(batch_id: int | str, min_num: int | None = None, timeout: float | None = None, clear_timeout_tasks: bool = True, return_partial_tasks: bool = False) → Tuple[List[Status], List[bytes]]

Wait for one batch and return task statuses plus serialized payload chunks.

async get_statuses(batch_id: int | str, min_num: int | None = None, timeout: float | None = None, clear_timeout_tasks: bool = True, return_partial_tasks: bool = False) → List[Status]

Wait for one batch and return only task statuses without materializing experiences.

async abort_batch(batch_id: int | str, return_partial_tasks: bool = False, restart_runners: bool = True) → None

Abort one batch and cleanup unfinished scheduler state.

has_step(batch_id: int | str) → bool

get_key_state(key: str) → Dict

Get the scheduler state.

Parameters:

key (str) – The key of the state to get.

Returns:

A dictionary of runner ids to their state for the given key.

Return type:

Dict

get_runner_state(runner_id: int) → Dict

Get the scheduler state.

Parameters:

runner_id (int) – The id of the runner.

Returns:

The state of the runner.

Return type:

Dict

get_all_state() → Dict

Get all runners’ state.

Returns:

The state of all runners.

Return type:

Dict

print_all_state() → None

Print all runners’ state in a clear, aligned table format.