Class Scheduler

trinity/explorer/scheduler.py:221–871

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

class Scheduler:
    """Scheduler for rollout tasks.

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

    def __init__(
        self,
        config: Config,
    ):
        self.logger = get_logger(__name__)
        self.config = config
        self.namespace = ray.get_runtime_context().namespace
        self.default_timeout = config.explorer.max_timeout * (config.explorer.max_retry_times + 1)
        self.max_retry_times = config.explorer.max_retry_times
        self.max_repeat_times = config.explorer.max_repeat_times_per_runner
        self.default_batch_size = config.buffer.batch_size
        self.running = False

        self.runner_num = (
            config.explorer.rollout_model.engine_num * config.explorer.runner_per_model
        )
        self.runners: Dict[int, RunnerWrapper] = dict()
        self.idle_runners: set[int] = set()  # runner_id of idle runners
        self.busy_runners: Dict[int, RunningTaskState] = dict()  # runner_id -> running state

        self.pending_tasks: Dict[Union[int, str], deque] = defaultdict(
            deque
        )  # batch_id -> (task, repeat_times, run_id_base)
        self.running_tasks: Dict[Union[int, str], set[asyncio.Future]] = defaultdict(
            set
        )  # batch_id -> futures
        self.task_num_map: Dict[Union[int, str], int] = defaultdict(
            int
        )  # batch_id -> tasks scheduled under this batch_id
        self.running_task_state_map: Dict[asyncio.Future, RunningTaskState] = dict()
        self.batch_is_eval_map: Dict[Union[int, str], bool] = dict()
        self.completed_tasks: Dict[
            Union[int, str], Dict[Union[int, str], CompletedTaskResult]
        ] = defaultdict(
            dict
        )  # batch_id -> results
        self.background_tasks: set[asyncio.Task] = set()

        self.scheduler_task: Optional[asyncio.Task] = None
        self.monitor_task: Optional[asyncio.Task] = None

        self.total_running_time = 0.0
        self.total_completed_steps = 0
        self.total_completed_sub_tasks = 0
        self.total_completed_tasks = 0

    async def _create_runner(
        self,
        runner_id: int,
    ):
        runner = RunnerWrapper(