迭代调度器

迭代调度器

本节摘要:无调度器的研究循环是带妄想的队列。调度器是循环决定停止探索什么的地方,那个决定就是整个博弈。把研究工作流建模成假设队列喂并行实验槽、结果扇回。用 asyncio 跑多实验并发使调度器能保所有槽忙。用 UCB 给每假设分支打分,使调度器能剪低产分支而不弃探索。把完成结果扇出到论文写作阶段与重排队列阶段,使高产分支产后续假设。露带分支分、槽占用、剪枝决策的每迭代 trace。平工作表按提交序跑作业,这在每作业独立时行,研究不独立——实验三的发现改实验四五的优先级,读结果扇入重排队列的调度器每单位算力做更有用功。有趣设计选择是评分规则:贪心总选当前领先者永不探索,均匀永不利用,UCB(upper confidence bound)是中道——利用领先者同时为试得少的分支留产能。

对应原课程:Phase 19 · Lesson 56 · iteration-scheduler(原英文 phases/19-capstone-projects/56-iteration-scheduler/docs/en.md)。本节属「AI Scientist」赛道第七节。

学习目标

阅读完本节,你应当能够:

1. 把研究工作流建模成假设队列喂并行实验槽、结果扇回。
2. 用 asyncio 跑多实验并发,使调度器保所有槽忙。
3. 用 UCB 给每假设分支打分,使调度器剪低产分支而不弃探索。
4. 把完成结果扇出到论文写作阶段与重排队列阶段,使高产分支产后续假设。
5. 露带分支分、槽占用、剪枝决策的每迭代 trace。

一、问题与直觉

为什么调度器而非工作表

平工作表按提交序跑作业,每作业独立时行。研究不独立:实验三的发现改实验四五的优先级。读结果扇入重排队列的调度器每单位算力做更有用功。有趣设计选择是评分规则:贪心总选当前领先者永不探索,均匀永不利用,UCB 是中道——利用领先者同时为试得少的分支留产能。

队列持假设,调度器在槽空时选最高 UCB 假设,每槽异步跑实验,完成实验把结果扇到总线,总线更新源分支的 UCB 统计、在分支产超阈值时扇出到论文写作阶段。

Hypothesis 的形状

UCB 评分

本节用的是经典 UCB1:

ucb(branch) = mean_reward(branch) + c * sqrt( ln(total_runs) / runs(branch) )
​

total_runs 是跨所有分支完成实验的总数。c 是探索权重,默认 sqrt(2)。零运行的分支得 +inf,使未试分支总先排。高均奖励分支保持高分直到其他分支追上;跑多次无多少奖励的分支被少跑的替代遮蔽。剪枝门与选择器分开:分支均奖励在至少 prune_after_runs(默认 3)次试验后落绝对地板(默认 0.2)下则剪,使队列有界。

asyncio 并行槽

调度器用 asyncio.create_task 驱实验。每任务跑实验运行器(async def 可调用对象)返 Result。主循环用 asyncio.wait(..., return_when=FIRST_COMPLETED) 等在飞任务集,每完成发评分更新。三槽并发,主循环永不阻在单实验。调度器在槽一空就起新任务,直到队列空且无在飞任务。

扇出:论文触发

分支均奖励跨 paper_threshold(默认 0.7)且该分支未产论文时,调度器把 paper.trigger 事件扇到输出列。下游第 52 节的论文写作器会拾起。本节触发捕为列使测试可断言。

扇出:后续假设

高产结果落时,调度器可调用户供的 expander 产同分支一个或多个后续假设。expander 是从 Result 到 list[Hypothesis] 的纯函数。本节发确定性 expander,对奖励超论文阈值的任何结果产两个后续。

预算

两预算保调度器不脱缰:max_experiments(跨所有分支跑的总实验数)与 max_seconds(墙钟上限,asyncio 时间)。任一触发,调度器停排新任务、等在飞的、返终 trace(含 stop_reason)。

二、从零实现

code/main.py 定义 Hypothesis、Result、BranchStats、IterationScheduler、make_deterministic_runner 工厂(返带可预测奖励的 asyncio 实验运行器,睡固定 delay_ms 默认 5ms 使并发可观察)。

class IterationScheduler:
    def __init__(self, queue, runner, slots=3, c=math.sqrt(2),
                 paper_threshold=0.7, prune_floor=0.2, prune_after=3):
        self.branches = defaultdict(BranchStats); self.paper_triggers = []

    async def run(self, max_experiments=20, max_seconds=30):
        trace = []; in_flight = set()
        while self.queue or in_flight:
            if len(in_flight) < self.slots and self.queue:
                hyp = self.queue.pop(self._pick_ucb())        # 最高 UCB
                in_flight.add(asyncio.create_task(self.runner(hyp)))
            done, in_flight = await asyncio.wait(
                in_flight, return_when=asyncio.FIRST_COMPLETED)
            for t in done:
                res = t.result()
                self._update_ucb(res)                          # 扇入更新统计
                if self._should_trigger(res): self.paper_triggers.append(res.branch)
                if self._should_prune(res.branch): self.queue.drop(res.branch)
                self.queue.extend(self.expander(res))          # 高产扇出后续
                trace.append(...)
        return SchedulerReport(trace, self.branches, self.paper_triggers, stop_reason)
​

code/tests/test_scheduler.py 覆盖:UCB 先选未试分支、并行槽占用、跨阈值时论文触发、低产试验后分支剪枝、后续假设扇出、预算退出(实验数与墙时两者)。

设计要点:UCB 是利用与探索的中道——mean_reward 利用领先者,c*sqrt(ln(total)/runs) 项为少跑分支留分。剪枝门与选择器分开:选择器用 UCB 平衡,剪枝用绝对地板硬剪,两者职责不同。扇出是研究的本质——高产分支产后续假设,使循环自增殖而非只消耗队列。三预算(队列空、实验数、墙时)保调度器不脱缰,stop_reason 露为何停。

三、框架对比

UCB1 是多臂老虎机经典算法(Auer 2002),AlphaGo 的蒙特卡洛树搜索、推荐系统的探索-利用、超参优化的 Hyperband 都用它或变体。Ray Tune、Optuna 的调度器在超参搜索上用类似带探索的评分。本节把 UCB 用于研究假设分支调度,是把bandit框架迁移到科研工作流。真实自主研究代理(AI-Scientist v2、MLAgentBench)的调度更复杂:多目标(标量奖励替成 Pareto 向量)、上下文bandit(按假设特征条件化)、持久统计(跨会话 checkpoint)。本节的 IterationScheduler 是这些扩展的共同骨架,asyncio 并行槽是产品化的起点(Ray/Celery 分布式替 asyncio)。

四、可复用产物

code/main.py + code/tests/test_scheduler.py。IterationScheduler.run(...) 是编排器入口,返 SchedulerReport(trace、每分支统计、论文触发、stop_reason)。报告供第 55 节端到端 demo 驱论文写作器。expander 是纯函数可独立复用——任何「高产结果产后续」场景。换真实验运行器(真实验执行),调度器与 UCB 逻辑不动。

五、练习

1. UCB vs 贪心:把 UCB 选择器换成纯贪心(只选最高均奖励),对比探索的分支数与最终最佳奖励。
2. 多槽扫描:扫槽数 1/2/4/8,绘墙时间与吞吐,确认收益递减。
3. 持久统计:给 UCB 统计加 checkpoint(存 JSON),重启后恢复探索预算。
4. 多目标:把标量奖励替成 Pareto 向量(准确率 + 速度),UCB 成 Pareto 选择器。
5. 上下文 bandit:让选择器按假设特征(长度、复杂度)条件化,使相似假设共享探索。

本节要点回顾

1. 调度器非工作表:研究不独立,读结果扇入重排队列的调度器每单位算力做更有用功。
2. UCB 中道:mean_reward(利用)+ c*sqrt(ln(total)/runs)(探索),零运行分支得 +inf 先排。
3. 剪枝门分离:UCB 平衡,绝对地板硬剪,两者职责不同。
4. asyncio 并行槽:wait(FIRST_COMPLETED),槽一空起新任务,永不阻单实验。
5. 扇出是本质:论文触发(超阈值)+ 后续假设(expander),使循环自增殖。
6. 三预算:队列空、实验数、墙时,stop_reason 露为何停。

下一节,我们做「端到端研究 demo」——把第 48~54 节的原语通过纯 Python import 组合,跑通种子→实验→调度→批判→论文的自终止循环。