TAI BUI
Back to curriculum
CapstonePython

Iteration Scheduler

A research loop without a scheduler is a queue with delusions. The scheduler decides what the loop stops exploring, and that decision is the whole game.

TIME: ~90 minutesPREREQUISITES: Phase 19 lessons 50-53

Learning Objectives

  • Model a research workflow as a hypothesis queue feeding parallel experiment slots whose results fan back in.
  • Run multiple experiments concurrently with asyncio so the scheduler can keep all slots busy.
  • Score each hypothesis branch with UCB so the scheduler can prune low-yield branches without abandoning exploration.
  • Fan out finished results to a paper-write stage and a re-queue stage so a high-yield branch spawns follow-up hypotheses.
  • Surface a per-iteration trace with slot occupancy, per-branch mean reward and run counts, and pruning decisions.

Why a scheduler, not a worklist

A worklist runs jobs in submission order, which is fine when each job is independent. A scheduler, as I use the word here, reads each result as it lands and reorders the queue. Research needs that, because a finding from experiment three changes the priority of experiments four and five, and a scheduler that reads the result fan-in gets more useful work done per unit of compute.

The design choice that matters is the scoring rule. A greedy scorer picks the current leader and never explores. A uniform scorer never exploits. UCB (upper confidence bound) sits between them: it exploits the leader while reserving capacity for branches that have been tried less.

Check: Why a scheduler, not a worklist0/1
1. What separates a scheduler, as this lesson uses the word, from a plain worklist?Define

The system shape

Mermaid diagram

The queue holds hypotheses. When a slot frees, the scheduler picks the highest-UCB hypothesis and runs it there as an asynchronous experiment. Finished experiments put their result onto the bus. The bus updates UCB statistics on the originating branch, and when a branch's yield crosses a threshold it also fans out to the paper-write stage.

Check: The system shape0/2
1. In this scheduler, what is a slot?Define
2. Why does the scheduler update a branch's runs and reward_sum as each result lands?Explain

The Hypothesis shape

Mermaid diagram

A hypothesis is one trial within a research direction, and branch names that direction. The Hypothesis record carries only id, branch, and payload. branch is the key for UCB statistics, and multiple hypotheses may share a branch. The statistics live on a separate BranchStats record per branch: runs is the count of completed experiments for that branch, and reward_sum is the cumulative reward. UCB reads both.

Check: The Hypothesis shape0/2
1. What is a hypothesis in this scheduler?Define
2. What does a branch's BranchStats record hold that UCB reads?Define

UCB scoring

UCB1 scores a branch as its mean reward plus an exploration bonus. This lesson uses the classic form.

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

total_runs is the count of all experiments completed across all branches. c is the exploration weight, and the lesson defaults it to sqrt(2). A branch with zero runs gets +inf, so untried branches are scheduled first. A branch with high mean reward keeps a high score until other branches catch up, and a branch that runs many times without much reward gets eclipsed by less-run alternatives.

The pruning gate is separate from the picker. It removes a branch from future scheduling when its mean reward falls below an absolute floor (default 0.2) after at least prune_after_runs trials (default 3). That keeps the queue bounded.

Check: UCB scoring0/12
1. What score does a branch with zero completed runs get under UCB1 in this scheduler?Define
2. In the UCB formula, what is the exploration weight c?Define
3. What is the pruning gate?Define
4. Why is always picking the branch with the highest mean reward a poor picker?Explain
5. UCB already ranks branches. What does the separate pruning gate add?Explain
6. With c = 0, a branch has 5 runs and a reward_sum of 1.0, and 10 experiments have completed in total. What does ucb_score return?Apply
7. Branch A has 4 runs with mean 0.8 and branch B has 1 run with mean 0.5, so 5 experiments have completed and c = sqrt(2). Which branch's queued hypothesis does the picker choose?Apply
8. A branch just finished its fifth run with a reward_sum of 1.0. The prune floor is 0.2 and prune_after_runs is 5. What does the scheduler do after that result?Apply
9. One branch has 6 queued hypotheses, every result has reward 0.1, there is 1 slot, max_experiments is 10, prune_floor is 0.2 and prune_after_runs is 3. How does the run end?Apply
10. In the untried-branches test, six hypotheses (A, B, C, A2, B2, C2 on branches A, B, C) run with one slot and every reward is 0.6. Which branches are picked first?Apply
11. At the end of the demo, branch-c has a mean of about 0.14, under the 0.25 prune floor, yet it is not pruned. Which ingredient of the gate explains this?Analyze
12. The picker keeps choosing a high-mean branch while less-run branches starve. Which single change moves the scheduler toward exploration without touching pruning?Design

Parallel slots with asyncio

The scheduler drives experiments with asyncio.create_task. Each task runs the experiment runner (an async def callable) and returns a Result. The main loop waits on the set of in-flight tasks with asyncio.wait(..., return_when=asyncio.FIRST_COMPLETED) and fires the scoring update on each completion.

Mermaid diagram

Three slots run concurrently, and the main loop does not block on any single experiment. It starts a new task as soon as a slot frees, and it stops once the queue is empty and no tasks are in flight.

Check: Parallel slots with asyncio0/3
1. How does the scheduler keep several slots busy at once?Explain
2. In the test of parallel slots, six hypotheses on different branches each sleep 0.05 s, run once with slots=1 and once with slots=3. Which outcome does the test assert?Apply
3. A teammate replaces asyncio.wait(FIRST_COMPLETED) with awaiting each in-flight task in dispatch order. What is the main consequence?Analyze

Fan-out: paper triggers

A paper trigger is a paper.trigger event. The scheduler fans one onto an output list when a branch's mean reward reaches paper_threshold (default 0.7) and that branch has not yet produced a paper. There is no minimum run count: after one run the mean is that run's reward, so a single lucky result can trigger a paper. Downstream, the paper writer from lesson fifty-four would pick it up. Here the trigger is captured as a list so tests can assert it.

Check: Fan-out: paper triggers0/3
1. When does the scheduler emit a paper.trigger for a branch?Define
2. In the once-per-branch test, four hypotheses on branch b1 each return reward 0.9 with paper_threshold 0.5 and one slot. What is report.paper_triggers?Apply
3. Under the lesson's rule that a trigger fires once a branch's mean reaches paper_threshold, a branch's first result is 0.95 with paper_threshold 0.7, and a paper trigger fires after that single run. Which part of the rule explains why one lucky run is enough?Analyze

Fan-out: follow-up hypotheses

A follow-up hypothesis is new work on the same branch, produced, when a high-yield result lands, by the user-supplied expander if one is set. The expander is a pure function from Result to list[Hypothesis]. The scheduler calls it only when an expander is set and the result's reward is at or above paper_threshold. The lesson ships a deterministic expander that always returns two follow-ups, so every qualifying result queues two new hypotheses, recorded as an expand event.

Check: Fan-out: follow-up hypotheses0/3
1. What is a follow-up hypothesis?Define
2. An expander is set and paper_threshold is 0.7. A result lands with reward exactly 0.7. What happens?Apply
3. In the demo, branch-a ran 17 times while branches b, c and d ran once each, though each was seeded with one hypothesis. What explains it?Analyze

Budgets

Two budgets protect the scheduler from runaway loops.

text
max_experiments    : total count of experiments dispatched across all branches
max_seconds        : wall-clock cap (monotonic time)

The scheduler stops dispatching new tasks as soon as either budget is hit. It also checks both after each batch of completions, and when one fires it stops the main loop, awaits the in-flight tasks, and returns a SchedulerReport. Its stop_reason field is "max_experiments" when the experiment budget stopped the run, "deadline" when the wall-clock cap did, and "queue_empty" when the queue and the in-flight set both ran dry. experiments_run counts dispatched tasks, including any still in flight when the budget hit. Those in-flight tasks are drained: their results update branch stats and can still fire a paper trigger, but no pruning or expansion runs for them.

Check: Budgets0/5
1. Which values can the stop_reason field of SchedulerReport take?Define
2. How do the two budgets differ?Explain
3. The max_experiments budget is hit while two tasks are still in flight. What happens to them?Apply
4. In the deadline test, 100 hypotheses are queued, each takes 0.05 s, there are 2 slots, max_experiments is 1000 and max_seconds is 0.12. What is the stop_reason?Apply
5. You launch a first run with an expander enabled and no idea how slow the runner is. Which configuration best protects you from runaway cost?Design

The Trace and final report

The trace is the ordered event log of the run. Each scheduling decision emits an event: dispatch (with the slot count after dispatch), result (with the branch mean and run count), prune, error for a failed experiment, and result.drain for results collected after a budget stop. A drain event carries only the branch and the reward, with no branch mean or run count, and when a drained result fires a paper trigger, the paper.trigger event is logged just before its result.drain event. A drained task that raised is skipped without an error event. Fan-out is two events: paper.trigger and expand. The final report summarises per-branch stats, experiments_run, wall_seconds, the stop_reason, and the paper triggers fired. The next lesson, the end-to-end demo, reads this report to drive the paper writer.

Check: The Trace and final report0/3
1. What is the trace in a SchedulerReport?Define
2. Per the lesson's description of the trace, why would it contain result.drain events instead of result events for some experiments?Explain
3. Two hypotheses on branch A are dispatched. The runner raises an exception for one and returns reward 0.5 for the other. What does the final report show?Apply
Author's Note

I'd keep the pruning gate separate from the picker. UCB ranks the queue, and pruning decides when a branch stops being a candidate at all. I'd rather read those as two rules than one blended score.

My bet is that prune_after_runs is the knob to tune first, because the floor of 0.2 only applies once a branch has at least prune_after_runs runs (default 3) behind its mean. I'd also set both budgets, max_experiments and max_seconds, from the first run, so a runaway loop is capped before it costs anything.

How to read the code

code/main.py defines Hypothesis, Result, BranchStats, IterationScheduler, and a make_deterministic_runner factory that returns an asyncio experiment runner with predictable rewards. The runner sleeps for a fixed delay_ms (default 5ms), which makes the concurrency observable.

code/tests/test_scheduler.py covers: UCB picks untried branches first, parallel slot occupancy, paper triggers when threshold is crossed, branch pruning after low-yield trials, fan-out follow-up hypotheses, and budget exit (both experiment count and wall clock).

Going further

Three extensions a real implementation will want. First, persistent UCB stats across sessions: the current statistics live in memory, and a real scheduler would checkpoint them so a restart preserves the exploration budget already spent. Second, multi-objective scoring: instead of a scalar reward, each result emits a vector and UCB becomes a Pareto-style picker. Third, contextual bandits: the picker conditions on hypothesis features (length, complexity) so similar hypotheses share exploration.

The scheduler is where research becomes more than a worklist. Once UCB is wired and the slots run in parallel, other improvements compose on top.

Warm-up0/3
1. Three jobs each spend most of their time waiting on a network call. Which approach in a single Python process finishes all three soonest?Warm-up
2. A team funds only whichever research direction looks best so far and never tries a new one. What is the main risk?Warm-up
3. An automated loop spawns follow-up work from its own results. Why cap the total number of runs and the elapsed time?Warm-up
Check: Going further0/1
1. A nightly restart wipes the scheduler's memory, so it re-explores branches it already knows are weak. Which extension fits the problem?Design

Code