Scheduler

The search engine beneath Search: the shared tree, branching gates, and allocation sessions.

Observations

Keep the evidence used to allocate continuations at one checkpoint: completed and pending purchases, plus running statistics of completed suffix returns. Evidence counts can differ from purchase counts because earlier rollouts also supply observations.

from focus.allocation import Observations

Definition

@dataclass
class Observations

Fields

m_done: int = 0
m_pending: int = 0
returns: RunningStats = field(default_factory=RunningStats)

Observations.m_alloc

@property
def m_alloc(self) -> int

Observations.n_eff

@property
def n_eff(self) -> int

AllocationLedger

Track completed and pending continuations at each node. Allocation sessions for the same tree share these counts.

from focus.scheduler import AllocationLedger

Definition

class AllocationLedger[P: Payload]

AllocationLedger.init

def __init__(self, *, observations: dict[int, Observations] | None=None) -> None

AllocationLedger.bind

def bind(self, view: SchedulerView[P]) -> None

AllocationLedger.counts

def counts(self, node: Node[P]) -> Observations

AllocationLedger.record

def record(self, parent: Node[P], chain: list[Node[P]]) -> None

Hook for recording the evidence from a completed expansion. The base implementation does nothing; RolloutLedger overrides it to record returns from continuation suffixes.

AllocationSession

Choose checkpoints and track their continuations within one pass. Implement choose() for a custom selection rule; the shared lifecycle handles completion, failure, and cancellation.

from focus.scheduler import AllocationSession

Definition

class AllocationSession[P: Payload](ABC)

AllocationSession.init

def __init__(self, ledger: AllocationLedger[P] | None=None) -> None

AllocationSession.start

def start(self, view: SchedulerView[P]) -> None

AllocationSession.stop

def stop(self) -> None

AllocationSession.choose

Required implementation.

@abstractmethod
def choose(self, view: SchedulerView[P]) -> Node[P] | None

AllocationSession.select

@final
def select(self, view: SchedulerView[P]) -> Reservation[P] | None

AllocationSession.active

def active(self, reservation: Reservation[P]) -> bool

AllocationSession.complete

@final
def complete(self, reservation: Reservation[P], state: SchedulerState[P], chain: list[Node[P]]) -> None

AllocationSession.fail

@final
def fail(self, reservation: Reservation[P], error: BaseException) -> None

AllocationSession.cancel

@final
def cancel(self, reservation: Reservation[P]) -> None

AllocationSession.selected

def selected(self, reservation: Reservation[P], view: SchedulerView[P]) -> None

AllocationSession.completed

def completed(self, reservation: Reservation[P], state: SchedulerState[P], chain: list[Node[P]]) -> None

AllocationSession.released

def released(self, reservation: Reservation[P], error: BaseException | None) -> None

Budget

Stop launching continuations when the configured turn, rollout, node, or goal limit is reached.

from focus.scheduler import Budget

Definition

class Budget[P: Payload](Termination[P])

Budget.init

def __init__(self, max_rollouts: int | None=None, *, max_turns: int | None=None, max_nodes: int | None=None, goal: Callable[[SchedulerView[P]], bool] | None=None) -> None

Budget.turns_spent

@staticmethod
def turns_spent(view: SchedulerView[P]) -> float

Turns realized, plus one per failure.

A dead expansion records no turns — it produced no payload to read them off — so a turn cap counting successes alone would never deplete against an expander that always fails, and the loop would re-select and re-fail forever. Charging a turn for the attempt is the same protection max_rollouts gets by counting failures, and it is a floor: a rollout that died partway spent at least that.

Budget.remaining

def remaining(self, view: SchedulerView[P]) -> int | None

ConcurrencyGating

Limit simultaneous continuations from each checkpoint while allowing more continuations after earlier ones finish.

from focus.scheduler import ConcurrencyGating

gating = ConcurrencyGating(limit=2)

A checkpoint can have two active continuations. A completed continuation releases its slot for another one.

Definition

class ConcurrencyGating(GatingPolicy)

ConcurrencyGating.init

def __init__(self, limit: int=1) -> None

ConcurrencyGating.gate

def gate(self, node: Node[Any]) -> Gate

Expander

Define how one selected node produces a chain of child payloads. This is the scheduler’s interface to task execution.

from focus.scheduler import Expander

Definition

class Expander[P](ABC)

Expander.expand

Required implementation.

@abstractmethod
async def expand(self, payload: P) -> tuple[list[P], Mapping[str, object]]

The children, plus metadata to record on each node they become.

The metadata rides back with the payloads rather than being written by the expander itself, because the payloads exist before the nodes do. An expander with nothing to record answers {}.

FifoSession

Maintain the state for one FIFO pass and choose the first eligible checkpoint from the frontier.

from focus.scheduler import FifoSession

Definition

class FifoSession[P: Payload](AllocationSession[P])

FifoSession.choose

def choose(self, view: SchedulerView[P]) -> Node[P] | None

Gate

Reserve and release branch capacity for a node without blocking the scheduler.

from focus.scheduler import Gate

Definition

class Gate

Gate.init

def __init__(self, *, max_children: int | None=None, max_pending: int | None=None) -> None

Gate.in_use

@property
def in_use(self) -> int

Gate.available

def available(self, node: Node[Any]) -> bool

Whether a branch may launch from node right now.

Gate.reserve

def reserve(self) -> None

Claim a slot (on launch).

Gate.release

def release(self) -> None

Free a slot (on completion).

GatingPolicy

Decide how many expansions a checkpoint may launch. Implement gate to vary branch limits by checkpoint or depth.

from focus.scheduler import GatingPolicy

Definition

class GatingPolicy(ABC)

GatingPolicy.gate

Required implementation.

@abstractmethod
def gate(self, node: Node[Any]) -> Gate

Node

Hold one tree payload, its parent and children, critic predictions, and scheduling metadata.

from focus.scheduler import Node

Definition

@dataclass
class Node[P: Payload]

Fields

id: int
parent_id: int | None
depth: int
payload: P
children: list[int] = field(default_factory=list)
gate: Gate = field(init=False)
critic: float | None = None
value_context: str | None = None
metadata: dict[str, Any] = field(default_factory=dict)

NodeScorer

Batch a critic’s predictions for new checkpoints through the execution services. The collector creates one; a search scores each new checkpoint before it can be selected.

from focus.scheduler import NodeScorer

scorer = NodeScorer(critic, services=agent.services)

Predictions are stored on each node as its critic prior, separate from observed returns. The critic’s context() renders each checkpoint as input.

Definition

class NodeScorer[P: Payload]

NodeScorer.init

def __init__(self, critic: Critic, *, services: ExecutionServices | None=None) -> None

NodeScorer.score

async def score(self, payloads: Sequence[P]) -> list[ValuePrediction]

Predict values for payloads in one batched call, preserving order.

A critic that answers with the wrong number of predictions raises here, inside the callers’ guards, so the fault costs an estimate rather than the run.

NodeScorer.write

def write(self, nodes: Sequence[Node[P]], scores: Sequence[ValuePrediction]) -> None

File scores on nodes, pairwise.

PASS_KEY

The node-metadata key recording which search phase created a node. Preparation trains only on edges created by search.

from focus.scheduler import PASS_KEY

Definition

PASS_KEY = 'pass'

Payload

Define the state properties required by the tree scheduler. RolloutState supplies the harness implementation.

from focus.scheduler import Payload

Definition

class Payload(Protocol)

Payload.forkable

@property
def forkable(self) -> bool

May the scheduler branch (expand) at this node.

Payload.done

@property
def done(self) -> bool

Whether this payload is terminal.

Payload.reward

@property
def reward(self) -> float

The default ranking score (used by SchedulerState.best).

RankedSession

An allocation session that ranks checkpoints. Implement score(); the session selects the eligible checkpoint with the highest score. Most extensions implement Priority instead, which builds this session for you.

from focus.scheduler import RankedSession

Definition

class RankedSession[P: Payload](AllocationSession[P], ABC)

RankedSession.score

Required implementation.

@abstractmethod
def score(self, view: SchedulerView[P], node: Node[P]) -> float

RankedSession.eligible

def eligible(self, score: float) -> bool

RankedSession.choose

def choose(self, view: SchedulerView[P]) -> Node[P] | None

Reservation

Identify one launched continuation. Two reservations from the same checkpoint track two separate continuations still in progress.

from focus.scheduler import Reservation

Definition

@dataclass(frozen=True, eq=False)
class Reservation[P: Payload]

Fields

node: Node[P]

RolloutExpander

Fork an isolated rollout, run it forward, and return its checkpoint chain to the scheduler.

from focus.scheduler import RolloutExpander

Definition

class RolloutExpander[S](Expander[RolloutState[S]])

RolloutExpander.init

def __init__(self, agent: Agent[S, Any], *, sampling_params: SamplingParams | None=None, annotation: Mapping[str, object] | None=None) -> None

RolloutExpander.expand

async def expand(self, payload: RolloutState[S]) -> tuple[list[RolloutState[S]], Mapping[str, object]]

RolloutExpander.start

async def start(self, prompt: str, *, workspace: Workspace | None=None) -> RolloutState[S]

Convenience: build the search root from a prompt (keeps the Agent here).

Scheduler

Coordinate concurrent expansions, allocation sessions, branching gates, and critic passes over a tree.

from focus.scheduler import Scheduler

Definition

class Scheduler[P: Payload]

Scheduler.init

def __init__(self, *, expander: Expander[P], allocation: AllocationSession[P] | None=None, termination: Termination[P] | None=None, gating: GatingPolicy | None=None, scorer: NodeScorer[P] | None=None, max_concurrency: int=8) -> None

Scheduler.run

async def run(self, root: P) -> SchedulerState[P]

Search from a root payload until termination, and return the tree it grew.

Scheduler.run_on

async def run_on(self, state: SchedulerState[P]) -> SchedulerState[P]

Search over an existing (possibly shared) tree until termination.

Lets other code run concurrently over the same state under state.lock. Gating comes from state.gating (set when the state was built), not this scheduler’s own default.

Returns the very state it was handed. The tree is the result: its terminals and its best node are questions asked of it (SchedulerState.terminals, SchedulerState.best), and a record that answered them once at return froze them at the moment the run stopped — which is exactly wrong for the caller that runs a second pass over the same tree, as every later search pass does.

Termination ends the launching, not the run: once remaining() is 0 the loop drains — it keeps awaiting and integrating what is already in flight until nothing is, and only then returns (see the module docstring for why, and for the overshoot that buys). The in-flight expansions are cancelled only when the caller cancels, or when something escapes the loop.

SchedulerState

Store the search tree and the frontier of nodes currently available for expansion.

from focus.scheduler import SchedulerState

Definition

@dataclass
class SchedulerState[P: Payload]

Fields

nodes: dict[int, Node[P]]
root_id: int
gating: GatingPolicy
frontier: dict[int, int] = field(default_factory=dict)
stats: dict[str, float] = field(default_factory=dict)
failures: list[tuple[int, str]] = field(default_factory=list)
lock: asyncio.Lock = field(default_factory=asyncio.Lock)
_next_id: int = 1
_arrivals: int = 0

SchedulerState.root

@classmethod
def root(cls, payload: P, *, gating: GatingPolicy) -> SchedulerState[P]

A fresh state seeded with the initial payload as the tree root.

SchedulerState.selectable

def selectable(self, node: Node[P]) -> bool

Whether node may be expanded: forkable, live, and its gate has a slot.

SchedulerState.refresh

def refresh(self, node_id: int) -> None

Recompute one node’s frontier membership (idempotent, order-preserving).

SchedulerState.add_child

def add_child(self, parent_id: int, payload: P) -> Node[P]

Attach one payload as a child of parent_id and index it.

SchedulerState.attach_chain

def attach_chain(self, parent_id: int, chain: list[P]) -> list[Node[P]]

Link a rollout’s payloads as a descendant line under parent_id.

The first payload becomes a child of parent_id (a new branch); each later payload is a child of the previous (the step line). Returns the new nodes.

SchedulerState.regate

def regate(self, gating: GatingPolicy) -> None

Re-mint every gate under a new policy and rebuild the frontier.

Gates are minted once, when a node is added, so a node that filled its discovery width has no slot left — correct while discovery runs, and exactly wrong for a later pass that must launch from those same nodes. This is the whole pass-boundary mechanism the tree needs: one explicit transition, called between passes under lock, rather than per-pass gate bookkeeping smeared through the loop.

SchedulerState.record_failure

def record_failure(self, node_id: int, exc: BaseException) -> None

SchedulerState.terminals

def terminals(self) -> list[Node[P]]

Every node whose payload ended the episode naturally.

SchedulerState.best

def best(self) -> Node[P] | None

The highest-reward terminal node, or None if nothing finished.

Ranked by the payload’s own reward — the one score every payload has. A caller ranking by anything else (a critic channel, a custom key) is choosing among terminals and can say so in its own terms.

SchedulerState.describe

def describe(self, batch: Any=None) -> str

The tree, one node per line, with any actor weights from batch joined.

See focus.export.format_tree for the line format.

SchedulerView

Read the rollout tree, its eligible checkpoints, and collection statistics. Priorities and estimators read this view to inspect parent-child relationships.

from focus.scheduler import SchedulerView

view = SchedulerView(tree)
root = view.root
children = view.children(root)

Wrap the tree returned by Collector.search(). Here, children contains the root’s direct branches, while descendants iterates through its whole subtree.

Definition

@dataclass(frozen=True)
class SchedulerView[P: Payload]

Fields

_state: SchedulerState[P]

SchedulerView.stats

@property
def stats(self) -> Mapping[str, float]

SchedulerView.failures

@property
def failures(self) -> tuple[tuple[int, str], ...]

Return failed expansions as (node_id, repr(exception)) pairs. The returned tuple is a snapshot of the failures recorded so far.

SchedulerView.nodes

def nodes(self) -> list[Node[P]]

Every node in the tree.

SchedulerView.selectable

def selectable(self, node: Node[P]) -> bool

SchedulerView.frontier_size

@property
def frontier_size(self) -> int

SchedulerView.iter_frontier

def iter_frontier(self) -> Iterator[Node[P]]

Iterate over eligible nodes without copying the frontier. Hold the tree lock while using this iterator.

SchedulerView.frontier

def frontier(self) -> list[Node[P]]

Eligible nodes in insertion order: each is live, forkable, and has an available gate slot.

SchedulerView.arrival

def arrival(self, node: Node[P]) -> int

When node last became selectable, as a position in arrival order.

SchedulerView.node

def node(self, node_id: int) -> Node[P] | None

Look up a node by ID, returning None if it is absent.

SchedulerView.root

@property
def root(self) -> Node[P]

The tree’s root node, which holds the starting state for this prompt.

SchedulerView.children

def children(self, node: Node[P]) -> list[Node[P]]

Return a node’s direct children in the order they were attached.

SchedulerView.terminals

def terminals(self) -> list[Node[P]]

Return nodes whose environment state has ended the episode.

SchedulerView.path_to

def path_to(self, node: Node[P]) -> list[Node[P]]

Return the parent chain ending at node, in root-to-node order. For a node in this tree, both root and node are included. A chain disconnected from this tree’s root produces a path that does not begin at the root.

SchedulerView.descendants

def descendants(self, node: Node[P]) -> Iterator[Node[P]]

Iterate depth-first over all nodes below node, excluding the node itself. The iterator is lazy, so callers can stop as soon as they find a match.

SchedulerView.len

def __len__(self) -> int

Termination

Define how many additional continuations the scheduler may launch from its current state.

from focus.scheduler import Termination

class NodeLimit(Termination):
    def __init__(self, limit):
        self.limit = limit

    def remaining(self, view):
        return 0 if len(list(view.nodes())) >= self.limit else 1

termination = NodeLimit(100)

The scheduler can launch one continuation while the tree has fewer than 100 nodes. A continuation can add several nodes, so this is an admission rule rather than an exact final tree-size bound.

Definition

class Termination[P: Payload](ABC)

Termination.remaining

Required implementation.

@abstractmethod
def remaining(self, view: SchedulerView[P]) -> int | None

Rollouts still launchable now: 0 = stop, positive = cap, None = none.

The scheduler counts in-flight rollouts against a returned cap, so a rollout budget is a hard cap — completed + in-flight never exceeds it.

0 ends the launching, not the run: the scheduler then drains what is already in flight and integrates it normally (finished rollouts are the expensive half, and cancelling them at the budget threw away up to one episode per slot). The first 0 is latched — this is never asked again on that run, so a predicate a drained result would flip back cannot re-open the pass — and the stats a caller reads afterwards include everything the drain integrated.

WidthGating

Cap the total number of direct child branches from each checkpoint, counting completed and in-flight branches together.

from focus.scheduler import WidthGating

gating = WidthGating(width=3)

Each checkpoint gets at most three direct branches. Completed children continue to count against this cap.

Definition

class WidthGating(GatingPolicy)

WidthGating.init

def __init__(self, width: int=1) -> None

WidthGating.gate

def gate(self, node: Node[Any]) -> Gate

ValuePrediction

Pair a critic score with the exact serialized context used to compute it.

from focus.scheduler.core.critic import ValuePrediction

Definition

@dataclass(frozen=True)
class ValuePrediction

Fields

score: float
context: str

token_entropies

Compute the per-token proxy -p log(p) from finite token log probabilities.

from focus.scheduler.core import token_entropies

Definition

def token_entropies(logprobs: Iterable[float]) -> list[float]

STAT_FAILURES

Name the scheduler statistic counting failed expansions.

from focus.scheduler.core.tree import STAT_FAILURES

Definition

STAT_FAILURES = 'failures'

STAT_NODES

Name the scheduler statistic counting tree nodes.

from focus.scheduler.core.tree import STAT_NODES

Definition

STAT_NODES = 'nodes'

STAT_ROLLOUTS

Name the scheduler statistic counting rollouts.

from focus.scheduler.core.tree import STAT_ROLLOUTS

Definition

STAT_ROLLOUTS = 'rollouts'

STAT_SCORE_FAILURES

Name the scheduler statistic counting critic-scoring failures.

from focus.scheduler.core.tree import STAT_SCORE_FAILURES

Definition

STAT_SCORE_FAILURES = 'score_failures'

STAT_TURNS

Name the scheduler statistic counting task transitions.

from focus.scheduler.core.tree import STAT_TURNS

Definition

STAT_TURNS = 'turns'

edge_duration

The number of task transitions on the edge into a node: the $d$ in the discount $\gamma^d$ that refined values and segment TD advantages apply to a successor’s value.

from focus.scheduler import edge_duration

Definition

def edge_duration(node: Node[AnyRollout], parent: Node[AnyRollout]) -> int

edge_return

Compute the discounted return along a node’s incoming edge.

from focus.scheduler import edge_return

Definition

def edge_return(node: Node[AnyRollout], parent: Node[AnyRollout], config: RewardConfig) -> float

finished_below

Enumerate completed trajectories reachable below a checkpoint.

from focus.scheduler import finished_below

Definition

def finished_below(view: SchedulerView[AnyRollout], node: Node[AnyRollout]) -> list[Node[AnyRollout]]

finished_through

Check whether a completed trajectory passes through a node, including the node itself.

from focus.scheduler import finished_through

Definition

def finished_through(view: SchedulerView[AnyRollout], node: Node[AnyRollout]) -> bool