From 6783fd72f85a1bce77b48c4c7b7f8b75baea8fbe Mon Sep 17 00:00:00 2001 From: Preetam Dwivedi Date: Fri, 24 Jul 2026 10:52:52 -0700 Subject: [PATCH] feat(speculation): generator contract and bestfirst impl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add submitqueue/extension/speculation/generator, the candidate-stream composition point (Generator/PathIterator) the standard Speculator pulls from, plus the bestfirst implementation and mocks. Open takes only the queue's batches. The generator offers every coherent path in the space, including paths whose builds already ran — suppressing finished paths belongs to the Allocator, the piece that already reconciles candidates against the stored path sets by ID. Keeping the generator ignorant of path sets leaves it a pure enumerator and removes the expand-before-skip subtlety from its iterator: Next pops, expands, returns, with no filtering loop. Ranking stays implementation-defined at the contract; only bestfirst commits to a scale. bestfirst ranks each path by how likely all its assumptions are to hold, and names paths long before it builds them. Each head roots a tree of its paths, and a single heap holds the frontier of every head's tree at once — as nodes, a node being a head plus the ascending set of alternatives that path takes rather than a built path. Handing a node out builds that one path and adds its one or two children. Pulling k paths therefore builds exactly k paths and sorts at most k heads' alternatives, however large the space behind them and however many heads the queue holds. Scores are summed log probabilities, which preserves ordering without underflow. Multiplying the probabilities instead loses the ordering outright: two heads waiting on 1200 and 1300 coin-flip dependencies both round to exactly 0.0, tie, and the heap falls back to its tie-break — so it can emit the 1300-dependency path first, one that is 2^100 times less probable. The ranking does not degrade, it inverts. Summing logs keeps the arithmetic in range at any depth, preserves the order because log is monotonic, and gives an assumption that cannot happen a score of negative infinity, which compares correctly against every finite score. A wide-head test pins this, asserting both that two 7100-dependency heads order correctly and that exponentiating either score yields zero. Scores from the injected scorer are validated on the way in. Everything downstream treats a score as a probability — its log is the ranking key, its complement is the other side's probability — so a value outside [0, 1] would not fail loudly, it would quietly produce a positive log, an inverted ordering, or a NaN that makes every comparison false. Each unresolved dependency has a preferred assumption — the one with the higher probability — and an unpreferred one. The head's best path takes preferred everywhere and roots the tree; every other node takes some subset of the head's alternatives, each adding its relativeScore. Subsets are walked by a canonical expansion under which every subset has exactly one parent, so each path is reached exactly once with no visited set, and a child never outranks its parent — which, since a node is the heap's maximum when it is popped, is what makes the emitted sequence non-increasing. A node's score is always summed from its head's best-path score in ascending alternative order, never carried over from the parent it grew out of. The two are not the same number: advancing {0,1} to {0,2} gives bestPathScore+r0+r2 summed against (s-r1)+r2 carried, and floating-point addition does not associate. Carrying would break the ordering guarantee silently and only on some inputs, so a test pins every emitted score to its canonical sum. Equal scores break on depth first, then on the alternatives taken, then on the head. Depth first offers every head's root before any head's deeper nodes, which matters because a dependency at probability exactly 0.5 deviates for free: ordering on the head instead would let one head's tied subtree drain a caller's build budget before another head's root was offered at all. Comparing the alternatives before the head then interleaves heads at equal depth, since position i is the i-th cheapest deviation for whichever head it belongs to. This replaces a tie-break on path ID, which needed a hash of every path the moment it entered the heap. Everything a dependency contributes — which assumption is preferred, what the preferred side scores, what deviating costs — follows from that batch alone, so it is worked out once per run rather than once per head waiting on it. Since entity.Batch.Dependencies holds a transitive closure, a chain of N batches has O(N^2) edges over N distinct dependencies, and the difference is not marginal: on a 3000-head chain Open drops from roughly 1.0s and 953MB to roughly 0.1s and 2MB. What Open cannot defer is the scoring itself, because ranking heads against each other needs every head's best-path score up front. Open holds the batches and their dependency slices rather than copying them, for as long as the iterator lives. The Generator contract now says so, since it constrains every caller and every implementation even though no signature changed. It is consistent with batches being replaced rather than edited in place, and it is what lets a head cost a struct rather than a slice per dependency. Generation honors context cancellation. Open checks before each head and newHead before each dependency, since a head can carry hundreds; Next checks before each pull, which is where a caller that has stopped consuming stops paying. All return the context error with no iterator and no candidate. Nothing after the heap pop can fail, so a node is never taken off the heap and then dropped: a pull that ends in an error leaves the stream exactly where it was, which a test pins. Both sides of each unresolved dependency are carried through as computed rather than derived from one another: 1-(1-score) is not score in floating point, and deriving one side would perturb every score below the cutoff. There is no depth bound. It existed to cap 2^n enumeration, and lazy generation removed that cost, so a head may now have as many unresolved dependencies as it likes. A dependency absent from the live batches cannot be scored, and gets a high default rather than an even split. It stays unresolved rather than becoming a certainty — both outcomes remain in the space. The cutoff used to pick which side is preferred is a separate constant from that default, so the two move independently; sharing one constant would have silently inverted the preferred assumption for every dependency below the default. Also sharpens scorer.Scorer's contract: it returns the probability that a batch's build succeeds, not "the likelihood of a successful land", which conflated the build with the merge. That is exactly the quantity bestfirst needs, since an assumption that a dependency succeeds is refuted when its builds cannot pass. The scorer README is rewritten around entity.Batch, which the interface has taken for some time. The package is three files by phase: bestfirst.go holds the generator, the per-dependency derivation and the scoring done at Open, head.go holds a head and everything derived from it including path materialization, and iterator.go holds the heap and the pull-time walk. --- submitqueue/extension/scorer/README.md | 63 +- submitqueue/extension/scorer/scorer.go | 14 +- .../speculation/generator/BUILD.bazel | 9 + .../extension/speculation/generator/README.md | 9 + .../generator/bestfirst/BUILD.bazel | 30 + .../speculation/generator/bestfirst/README.md | 100 ++ .../generator/bestfirst/bestfirst.go | 324 +++++ .../generator/bestfirst/bestfirst_test.go | 1052 +++++++++++++++++ .../speculation/generator/bestfirst/head.go | 230 ++++ .../generator/bestfirst/iterator.go | 195 +++ .../speculation/generator/generator.go | 57 + .../speculation/generator/mock/BUILD.bazel | 13 + .../generator/mock/generator_mock.go | 98 ++ 13 files changed, 2135 insertions(+), 59 deletions(-) create mode 100644 submitqueue/extension/speculation/generator/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/README.md create mode 100644 submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/bestfirst/README.md create mode 100644 submitqueue/extension/speculation/generator/bestfirst/bestfirst.go create mode 100644 submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go create mode 100644 submitqueue/extension/speculation/generator/bestfirst/head.go create mode 100644 submitqueue/extension/speculation/generator/bestfirst/iterator.go create mode 100644 submitqueue/extension/speculation/generator/generator.go create mode 100644 submitqueue/extension/speculation/generator/mock/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/mock/generator_mock.go diff --git a/submitqueue/extension/scorer/README.md b/submitqueue/extension/scorer/README.md index 21c4ee60..0a7a33ad 100644 --- a/submitqueue/extension/scorer/README.md +++ b/submitqueue/extension/scorer/README.md @@ -1,64 +1,17 @@ -# Scorer +# scorer -Vendor-agnostic interface for computing success probability scores for code changes. +A `Scorer` returns the probability that a batch's build succeeds — a number between 0.0 and 1.0. It is handed the batch identity and resolves the batch's changes itself through an injected `changeset.Resolver`, so callers pass an `entity.Batch` and nothing more. -## Interface +Callers may score every batch a queue is waiting on, so implementations should be cheap. A speculation run scores each batch at most once, but it does not carry results across runs; anything expensive belongs behind the implementation's own cache. -### Scorer - -Computes a success probability for a given change. - -```go -type Scorer interface { - Score(ctx context.Context, change entity.Change) (float64, error) -} -``` - -- **change**: A `entity.Change` identifying the code change to score. -- **Score**: Returns a probability between 0.0 and 1.0 indicating the likelihood of a successful land. Returns an error if scoring fails. +Like the other extensions, a `Scorer` is selected **per queue** by the wiring layer through the `Config` (queue name) and `Factory` interface. ## Implementations -### Heuristic - -Scores a change by extracting a numeric value via a `ValueFunc` and matching it against ordered buckets. Each bucket maps a `[Min, Max]` range to a probability. - -```go -s := heuristic.New( - []heuristic.Bucket{ - {Min: 0, Max: 5, Score: 0.95}, - {Min: 6, Max: 20, Score: 0.75}, - {Min: 21, Max: 100, Score: 0.5}, - }, - func(ctx context.Context, change entity.Change) (int, error) { - // resolve the change into a numeric metric - return filesChanged, nil - }, -) - -score, err := s.Score(ctx, change) -``` - -### Composite - -Combines multiple named scorers into a single score using a reduce function. The reduce function receives a `map[string]float64` mapping scorer names to their scores, enabling domain-aware aggregation. - -Built-in reduce functions: `Min`, `Max`, `Avg`. - -```go -s := composite.New( - map[string]scorer.Scorer{ - "files": fileScorer, - "deps": depScorer, - }, - composite.Min, -) +**`heuristic`** scores a batch by extracting one number from its changes and matching that against ordered buckets, each mapping a `[Min, Max]` range to a probability. The extraction is a caller-supplied `ValueFunc` over the resolved `entity.BatchChanges`, so the same bucketing works for files touched, lines changed, or any other metric. -score, err := s.Score(ctx, change) -``` +**`composite`** runs several named scorers and reduces their scores to one. The reduce function receives the scores keyed by scorer name, so it can weigh sources differently rather than treating them as interchangeable; `Min`, `Max`, and `Avg` are provided. -## Implementing a Backend +## Adding a backend -1. Create `extension/scorer/{backend}/` directory -2. Implement the `Scorer` interface -3. Accept `entity.Change` and resolve it into whatever data the implementation needs +Create a package under `scorer//` whose `New(...)` returns a `scorer.Scorer`, injecting whatever it needs at construction — a `changeset.Resolver` to reach the batch's changes, a metrics scope, any client. Do not add a `Config` or `Factory` implementation here; per-queue routing and the factory adapter live in the wiring layer. diff --git a/submitqueue/extension/scorer/scorer.go b/submitqueue/extension/scorer/scorer.go index b3af1b09..da5ce729 100644 --- a/submitqueue/extension/scorer/scorer.go +++ b/submitqueue/extension/scorer/scorer.go @@ -22,11 +22,17 @@ import ( "github.com/uber/submitqueue/submitqueue/entity" ) -// Scorer computes a success probability score for a batch based on its changes. +// Scorer computes the probability that a batch's build succeeds, based on its +// changes. type Scorer interface { - // Score returns a probability between 0.0 and 1.0 indicating the likelihood - // of a successful land for the given batch. It is handed the batch identity - // and resolves the batch's changes itself through an injected changeset.Resolver. + // Score returns a probability between 0.0 and 1.0 that the given batch's + // build succeeds. It is handed the batch identity and resolves the batch's + // changes itself through an injected changeset.Resolver. + // + // Callers may score every batch a queue is waiting on, so implementations + // should be cheap: a speculation run scores each batch at most once, but it + // does not carry results over to the next run, so anything expensive to + // compute belongs behind the implementation's own cache. Score(ctx context.Context, batch entity.Batch) (float64, error) } diff --git a/submitqueue/extension/speculation/generator/BUILD.bazel b/submitqueue/extension/speculation/generator/BUILD.bazel new file mode 100644 index 00000000..2e290235 --- /dev/null +++ b/submitqueue/extension/speculation/generator/BUILD.bazel @@ -0,0 +1,9 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["generator.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/generator", + visibility = ["//visibility:public"], + deps = ["//submitqueue/entity:go_default_library"], +) diff --git a/submitqueue/extension/speculation/generator/README.md b/submitqueue/extension/speculation/generator/README.md new file mode 100644 index 00000000..f44f5852 --- /dev/null +++ b/submitqueue/extension/speculation/generator/README.md @@ -0,0 +1,9 @@ +# generator + +The `generator` package is a piece the `standard` `Speculator` is built from: a `Generator` produces the queue's candidate paths as one stream, best first, across all heads. It is **not** controller-facing — the speculate controller only knows the `Speculator` contract, and a different `Speculator` need not split its work this way. So there is no `Config` or `Factory` here; a `Generator` is chosen when the `standard` `Speculator` is constructed. + +`Open` starts the stream over the queue's live batches and returns a `PathIterator`. The caller pulls one candidate at a time and the generator does only the work that answer needs. A cancelled or expired context ends the stream with its error. + +Candidates never repeat and never contradict a known fact. Beyond that, the order is the `Generator`'s own: it yields candidates in whatever ranking it implements, and each carries the score it ranked by — higher first, on a scale the generator defines. Consumers take the iterator in the order given and do not interpret the score. Scores mean something only within the run and are never stored. + +The generator offers every path in the space, including paths whose builds already ran. Suppressing finished paths is the `Allocator`'s job, since that is the piece reconciling candidates against the stored path sets. diff --git a/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel new file mode 100644 index 00000000..ebdca570 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel @@ -0,0 +1,30 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = [ + "bestfirst.go", + "head.go", + "iterator.go", + ], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/generator/bestfirst", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/scorer:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = ["bestfirst_test.go"], + embed = [":go_default_library"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/scorer:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/generator/bestfirst/README.md b/submitqueue/extension/speculation/generator/bestfirst/README.md new file mode 100644 index 00000000..5ada2bc5 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/README.md @@ -0,0 +1,100 @@ +# bestfirst + +A *head* — the batch we want to build — waiting on unfinished dependencies has one path per combination of outcomes — 2ⁿ of them — while callers want the best few. `bestfirst` hands them out in score order without building the rest. + +Dependencies that already resolved carry a forced assumption: succeeded means the head builds on top of it, failed or cancelled means it builds without. They stay in the path but drop out of the search. Each unresolved one is a two-way assumption whose *probability* — that the dependency's build succeeds — comes from an injected `scorer.Scorer`, remembered per dependency so one shared by several heads is scored once. A dependency that is not a live batch cannot be scored at all, and gets a high default, since nearly every batch a queue accepts does build successfully. A path's score follows from the probability that all its assumptions hold. + +**The best path is not "assume everything succeeds."** Each dependency takes its *preferred* assumption — the one with the higher probability — so a dependency that will probably fail is assumed to fail. Every other path takes the unpreferred assumption on some of them, and each of those costs a known factor. + +**Scores are logarithms.** A probability is the `[0, 1]` value the scorer gives; a score is always its logarithm. Scores use summed log probabilities to preserve ordering without underflow: a plain product across a wide head rounds to zero, ties, and loses the order entirely. A score is at most 0, and an assumption that cannot happen scores negative infinity. The package doc in `bestfirst.go` shows the failure it avoids. + +**Laziness.** Each head roots a tree of its paths: the root takes no alternatives, and every other node takes some subset. One heap holds the frontier of all those trees at once — the nodes reached but not yet handed out. Handing a node out adds its children, so the heap always holds each head's current best and the top of it is the best path in the queue. + +A node is its head, the alternatives it takes, and a score, so reaching one costs nothing near the size of the head. The path itself is built only when the node is handed out, and a head works out and sorts its alternatives only when its root is expanded — a head nobody pulls from does neither. Pulling *k* paths builds *k* paths and works out the alternatives of at most *k* heads, whatever the 2ⁿ behind them. + +What cannot be deferred is scoring. Ranking heads against each other needs every head's best-path score up front, so `Open` walks every dependency of every head once. That is linear in the queue's dependency edges — and since `Dependencies` holds a transitive closure, a chain of batches has quadratically many edges in its depth. `Open` pays that whether or not anything is pulled. + +`bestfirst` discards nothing: pull long enough and every path for every head comes out. It does no conflict relaxation (`ignored` assumptions) — that is the `Speculator`'s call, not the generator's. + +## Worked example + +A head waits on three unresolved dependencies. The scorer gives each the probability its build succeeds: + +| dependency | probability | preferred assumption | +| --- | --- | --- | +| `d0` | 0.9 | succeeds | +| `d1` | 0.8 | succeeds | +| `d2` | 0.6 | succeeds | + +All three are likelier to succeed than not, so the head's **best path** assumes all three succeed. Its probability is the product `0.9 × 0.8 × 0.6 = .432`, and its score is the sum of the logs, `−0.105 + −0.223 + −0.511 = −0.839`. Probabilities read more easily, so the tables below stay in them; the code only ever compares logs. + +### Step 1 — the alternatives, cheapest first + +Flipping one dependency to its other side swaps its factor in that product, multiplying the result by `unpreferred / preferred`. Sorted by what that costs: + +| | flips | ratio | log | +| --- | --- | --- | --- | +| `a0` | `d2` (0.6) | 0.4 / 0.6 = .67 | −0.41 | +| `a1` | `d1` (0.8) | 0.2 / 0.8 = .25 | −1.39 | +| `a2` | `d0` (0.9) | 0.1 / 0.9 = .11 | −2.20 | + +`a0` is `d2`, not `d0`. The cheapest deviation is from the *least* confident dependency — it is the assumption we mind changing least. + +This is the only sorting that happens, it is per head, and it happens the first time something pulls from that head. + +### Step 2 — the tree + +Every path is the best path with some set of alternatives taken. `expand` grows those sets one step at a time: either take the next alternative as well, or move the last one along. That lays all 2³ = 8 paths out as a tree, rooted at the best path: + +``` +{} .432 flip nothing — the best path +└── {a0} .288 flip d2 + ├── {a0,a1} .072 flip d2, d1 + │ ├── {a0,a1,a2} .008 flip all three — the worst path + │ └── {a0,a2} .032 flip d2, d0 + └── {a1} .108 flip d1 + ├── {a1,a2} .012 flip d1, d0 + └── {a2} .048 flip d0 +``` + +Two things to read off it. Every node scores at or below its parent — that is what the heap relies on, and it holds because the alternatives are sorted, so both moves reach for a later, never-cheaper one. But **siblings are not ordered**: `{a0,a1}` at .072 sits above `{a1}` at .108 in the tree while scoring below it. The tree only guarantees the parent relation. Global order is the heap's job. + +### Step 3 — scoring a node + +A node's score is the head's best-path score plus each taken alternative's log, and never anything carried over from its parent. For `{a0,a2}`: + +``` +−0.8393 best path +−0.4055 a0 +−2.1972 a2 +─────── +−3.4420 +``` + +Which is the same path read directly: `d0` fails, `d1` succeeds, `d2` fails → `0.1 × 0.8 × 0.4 = .032`, and `log(.032) = −3.4420`. + +### Step 4 — pulling + +`Open` puts just the root in the heap. Each pull takes the best node there, builds *that* path, and puts back its children: + +``` +handed out heap afterwards +{} .432 {a0}·.288 +{a0} .288 {a1}·.108 {a0,a1}·.072 +{a1} .108 {a0,a1}·.072 {a2}·.048 {a1,a2}·.012 +{a0,a1} .072 {a2}·.048 {a0,a2}·.032 {a1,a2}·.012 {a0,a1,a2}·.008 +{a2} .048 {a0,a2}·.032 {a1,a2}·.012 {a0,a1,a2}·.008 +{a0,a2} .032 {a1,a2}·.012 {a0,a1,a2}·.008 +{a1,a2} .012 {a0,a1,a2}·.008 +{a0,a1,a2} .008 — +``` + +The scores come out in descending order without the eight paths ever being enumerated or sorted together, and the heap never holds more than four nodes. Stop after two pulls and only two paths were ever built — everything else reached is still just a set of indexes. + +With several heads in the queue there is still only this one heap. Every head's root goes in at `Open`, they compete on score, and a head nobody reaches never even sorts its alternatives. + +### Ties + +Equal scores break on depth first, so every head's root is offered before any head's deeper nodes; then on the alternatives taken, which interleaves heads at equal depth; then on the head itself. Depth has to come first because a dependency at probability exactly 0.5 deviates for free, and ordering on the head instead would let one head's tied subtree drain a caller's budget before another head's root was offered at all. + +The proof that this walk reaches every path exactly once, in score order, is on `expand` in `iterator.go`. The behavior is pinned by tests in `bestfirst_test.go`. diff --git a/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go new file mode 100644 index 00000000..405fde01 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go @@ -0,0 +1,324 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package bestfirst hands out speculation paths for a queue, best ones first. +// +// A batch we want to build (a "head") often depends on other batches that have +// not finished yet. Each unresolved dependency could go either way, so there is +// more than one sensible thing to build: one path per combination of outcomes. +// +// The package's vocabulary: +// +// - A head is the batch we want to build, and the root of a tree of the paths +// it could be built along. +// - An assumption is what a path expects of one dependency — that it succeeds +// or that it does not. Once the dependency resolves the assumption is +// forced rather than chosen. Of the two, the one with the higher +// probability is preferred and the other unpreferred. +// - An alternative is the option to take the unpreferred assumption for a +// dependency still unresolved. The head's best path takes none of them; +// every other path takes some subset. +// - A node is one path waiting in the heap with its score, standing for it — +// a head plus the alternatives it takes — without building it. +// +// Throughout, a probability is the [0, 1] value a scorer gives; a score is +// always its logarithm. +// +// # How it works +// +// Open is handed the queue's live batches. For every batch still speculating: +// +// 1. Score each dependency it waits on. A dependency is scored once per run +// however many heads wait on it, and one whose build already finished is a +// fact rather than a probability, so it drops out of the search. +// 2. Total the head's best path — every remaining dependency on its likelier +// assumption. That total is the head's opening score. +// 3. Put that best path into a single heap shared by every head, as a node: +// the head, plus the (still empty) set of alternatives it takes. +// +// Then each pull takes the highest-scoring node in the heap, builds its path, +// and puts back the one or two nodes that come after it within that same head. +// So the heap always holds every head's current best, and its top is the best +// path in the queue. Nothing beyond those nodes is ever built, and no head's +// alternatives are worked out until something pulls from it. +// +// The README walks a three-dependency head through this end to end. +// +// # Scores are logarithms +// +// Scores use summed log probabilities to preserve ordering without underflow. +// Multiplying the probabilities instead loses the ordering outright: two heads +// waiting on 1200 and 1300 coin-flip dependencies both round to 0.0 (the +// smallest positive float64 is about 2^-1074), so they tie, the heap falls +// back to its tie-break, and it can emit the 1300-dependency path first — one +// that is 2^100 times less probable. +// +// A score is therefore at most 0, more negative meaning less probable, and +// negative infinity for an assumption that cannot happen. Scores combine by +// addition, never multiplication. Nothing exponentiates; the score is a +// ranking key, never a number anyone reports. +// +// # Laziness +// +// A head with n unresolved dependencies has 2^n paths, and callers want the +// best handful. A node costs nothing near n to hold, so the heap holds nodes +// and a path is built only when it is handed out. Pulling k paths builds k +// paths and works out the alternatives of at most k heads, whatever n is and +// however many heads the queue has. +// +// Scoring is what Open cannot defer: ranking heads against each other needs +// every head's best-path score up front. The README gives the cost that leaves. +package bestfirst + +import ( + "context" + "fmt" + "math" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/scorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +const ( + // probabilityCutoff is where the preferred assumption flips: at or above it + // a dependency is assumed to succeed, below it assumed to fail. + // + // Not a tuning knob. The preferred side is whichever one this favors, so + // moving it would leave that side below 0.5 and relativeScore above 0, + // which breaks the best-first ordering. + probabilityCutoff = 0.5 + + // defaultProbability is used for a dependency that is not among the live + // batches and so cannot be scored. Nearly every batch a queue accepts does + // build successfully, so a high default is closer to the truth than an even + // split and keeps an unscoreable dependency from dragging down every path + // that assumes it succeeds. + defaultProbability = 0.95 +) + +// gen builds best-first iterators. scorer gives each dependency's probability +// of building successfully. +type gen struct { + scorer scorer.Scorer +} + +// New returns a best-first generator.Generator. scorer scores each unresolved +// dependency. +func New(sc scorer.Scorer) generator.Generator { + return gen{scorer: sc} +} + +// Open scores every dependency, totals each head's best-path score, and queues +// the root of that head's tree. All the scoring happens here, which is why +// Next never calls the scorer. +// +// The batches and their dependency slices are held, not copied, for as long as +// the iterator lives; generator.Generator says what that asks of callers. +// +// A cancelled or expired ctx aborts with its error and no iterator. +func (g gen) Open(ctx context.Context, batches []entity.Batch) (generator.PathIterator, error) { + batchByID := make(map[string]entity.Batch, len(batches)) + for _, b := range batches { + batchByID[b.ID] = b + } + cache := newDependencyCache(g.scorer, batchByID) + + it := &iterator{} + for _, batch := range batches { + // Scoring a queue's worth of heads can take a while, so give up as soon + // as the caller stops waiting rather than at the end. + if err := ctx.Err(); err != nil { + return nil, err + } + // Only Speculating heads get paths. Batches in any other state (created, + // merging, cancelling, finished) still tell us how their dependencies + // turned out, but we never propose work on them. + if batch.State != entity.BatchStateSpeculating { + continue + } + h, err := newHead(ctx, batch, cache) + if err != nil { + return nil, err + } + // Queue the root of the head's tree: no alternatives taken, every + // dependency on its preferred assumption. Nothing is built for it yet — + // not the head's alternatives, not the path itself. + it.push(h, []int{}) + } + return it, nil +} + +// dependency is what one dependency batch contributes to any path over it: +// which assumption is preferred, what the preferred side scores, and what +// deviating costs. It follows from the batch alone, so heads that share a +// dependency share this, and its logarithms are taken once per queue rather +// than once per head that depends on it. +type dependency struct { + // unresolved reports whether the outcome is still open. A resolved + // dependency forces its assumption and offers no alternative, so the two + // unpreferred fields below mean nothing for it. + unresolved bool + // preferred is the assumption a path takes unless it deviates: for a + // resolved dependency the one its state forces, otherwise the likelier of + // the two. + preferred entity.DependencyAssumption + // unpreferred is the assumption a path takes by deviating. + unpreferred entity.DependencyAssumption + // preferredScore is what this dependency adds to a head's best-path score: + // the log of the preferred side's probability. Zero when resolved — a + // certainty costs nothing. Always finite, since a preferred probability is + // at least 0.5. + preferredScore float64 + // relativeScore is what deviating adds to a path's score. Always at most 0, + // since the unpreferred side never has the higher probability, and negative + // infinity when the unpreferred side cannot happen. + // + // A difference between the two sides rather than just the unpreferred + // side's log, because the best-path score already counts the preferred one: + // with preferred 0.8 and unpreferred 0.2, swapping 0.8 out of the sum for + // 0.2 is one add of log(0.2)-log(0.8). + relativeScore float64 +} + +// dependencyCache works each dependency out at most once per run. Several heads +// usually wait on the same dependency, and there is no reason to score it — or +// take its logarithms — once per head. +// +// The cache lives for one Open and no longer. A later run works the same batch +// out again, deliberately: a score can move as the queue does. scorer.Scorer's +// contract says the same and asks implementations to be cheap, so anything +// expensive to compute belongs behind the implementation's own cache, where it +// can decide what is safe to reuse across runs. +type dependencyCache struct { + scorer scorer.Scorer + batches map[string]entity.Batch + cached map[string]dependency +} + +func newDependencyCache(sc scorer.Scorer, batches map[string]entity.Batch) *dependencyCache { + return &dependencyCache{scorer: sc, batches: batches, cached: make(map[string]dependency)} +} + +// known returns what has already been worked out for the given dependency, and +// whether anything has. The zero value is not a valid answer, so a caller that +// might be looking one up for the first time must check ok. +// +// Deliberately just the map read, so it is cheap enough for the compiler to +// inline it: this runs once per dependency of every head, which on a deep queue +// is millions of times. Pairing it with resolve at the call site rather than +// wrapping the two in one method is what keeps it inlinable. +// +// Open works out every dependency of every speculating head, so a caller +// walking a head's own dependency list after that always hits. +func (c *dependencyCache) known(batchID string) (dependency, bool) { + d, ok := c.cached[batchID] + return d, ok +} + +// resolve works out a dependency known has not seen before, and remembers it. +// +// A score outside [0, 1] is rejected rather than used. Everything downstream +// treats it as a probability — its log is the ranking key, its complement is +// the other side's probability — so a bad value would not fail loudly. It +// would produce a positive log, an ordering that is no longer best-first, or a +// NaN that makes every comparison false. +func (c *dependencyCache) resolve(ctx context.Context, batchID string) (dependency, error) { + batch, live := c.batches[batchID] + // A resolved dependency is never scored: the scorer has nothing to say + // about a build that already finished. Only a live batch can be resolved — + // one that is not in the queue has no state to read, so it stays open. + if assumption, resolved := resolvedAssumption(batch.State); live && resolved { + // preferredScore stays 0, and 0 is log(1): a certainty multiplies the + // path's probability by 1, so a resolved dependency adds nothing to any + // head's total. unresolved stays false, so it offers no alternative and + // drops out of the search — the path takes the fact and moves on. + d := dependency{preferred: assumption} + c.cached[batchID] = d + return d, nil + } + + // Not a live batch, so there is nothing to score. Its outcome is still + // open, though, so it keeps both assumptions and just sits at the default. + score := defaultProbability + if live { + s, err := c.scorer.Score(ctx, batch) + if err != nil { + return dependency{}, err + } + if math.IsNaN(s) || s < 0 || s > 1 { + return dependency{}, fmt.Errorf("scorer returned %v for batch %q: want a probability in [0, 1]", s, batchID) + } + score = s + } + + // Prefer the assumption with the higher probability; on an exact tie, + // prefer succeeds. The preferred probability is therefore always at least + // 0.5, so its log is always finite. + // + // Both sides are carried through as computed: deriving one from the other + // is not free in floating point (1-(1-score) is not score — for score 0.1 + // it is 0.09999999999999998), and that would perturb every score below the + // cutoff. + preferredProbability, unpreferredProbability := score, 1-score + preferredAssumption, unpreferredAssumption := entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails + if score < probabilityCutoff { + preferredProbability, unpreferredProbability = 1-score, score + preferredAssumption, unpreferredAssumption = entity.DependencyAssumptionFails, entity.DependencyAssumptionSucceeds + } + + d := dependency{ + unresolved: true, + preferred: preferredAssumption, + unpreferred: unpreferredAssumption, + preferredScore: math.Log(preferredProbability), + // A path's probability is the product of its assumptions' + // probabilities, so deviating here swaps one factor for the other: + // the product is multiplied by unpreferred/preferred. Scores are the + // logs of those products, and log turns that multiplication into the + // addition of log(unpreferred/preferred) — the subtraction below. + // + // When the unpreferred side cannot happen its probability is 0, so the + // product goes to 0 and log(0) is -Inf. That is the right answer rather + // than an overflow: an impossible path sinks below every possible one, + // and -Inf stays -Inf however many finite terms are added to it. + relativeScore: math.Log(unpreferredProbability) - math.Log(preferredProbability), + } + c.cached[batchID] = d + return d, nil +} + +// resolvedAssumption reports the assumption a dependency's state forces, if +// that state is final. An unresolved dependency has no forced assumption and +// stays in the search. +// +// Only terminal states resolve. Cancelling deliberately does not: cancellation +// is best-effort, so a cancelling batch can still succeed. Resolving it to +// fails would drop the succeeds side out of the search altogether, and if the +// cancel then lost the race every path for the head would be refuted at once. +// Leaving it unresolved keeps both outcomes in the space, which is what an +// unresolved dependency means. +func resolvedAssumption(state entity.BatchState) (entity.DependencyAssumption, bool) { + switch state { + case entity.BatchStateSucceeded: + // Already succeeded, so building on top of it is a fact, not an + // assumption at risk. + return entity.DependencyAssumptionSucceeds, true + case entity.BatchStateFailed, entity.BatchStateCancelled: + // Never succeeding, so building without it is a fact too. + return entity.DependencyAssumptionFails, true + default: + return entity.DependencyAssumptionUnknown, false + } +} diff --git a/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go new file mode 100644 index 00000000..926e8dd2 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go @@ -0,0 +1,1052 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bestfirst + +import ( + "context" + "fmt" + "math" + "math/rand" + "sort" + "strings" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/scorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +// stubScorer scores each batch by ID, defaulting to 0.5 for unknown batches. It +// is a minimal scorer.Scorer for exercising the generator without a resolver. +type stubScorer struct { + scores map[string]float64 +} + +func (s stubScorer) Score(_ context.Context, b entity.Batch) (float64, error) { + if v, ok := s.scores[b.ID]; ok { + return v, nil + } + return 0.5, nil +} + +func scored(scores map[string]float64) scorer.Scorer { + return stubScorer{scores: scores} +} + +// drainAll pulls every candidate from an iterator. +func drainAll(t *testing.T, iter generator.PathIterator) []entity.CandidatePath { + t.Helper() + var out []entity.CandidatePath + for { + c, ok, err := iter.Next(context.Background()) + require.NoError(t, err) + if !ok { + break + } + out = append(out, c) + } + return out +} + +// forHead returns only the candidates whose head is headID. +func forHead(cands []entity.CandidatePath, headID string) []entity.CandidatePath { + var out []entity.CandidatePath + for _, c := range cands { + if c.Path.Head == headID { + out = append(out, c) + } + } + return out +} + +// assumptionFor returns what a path assumes about a given dependency batch. +func assumptionFor(p entity.SpeculationPath, dep string) entity.DependencyAssumption { + for _, d := range p.Dependencies { + if d.Batch == dep { + return d.Assumption + } + } + return entity.DependencyAssumptionUnknown +} + +// assumptionKey renders a path's assumptions in queue order, to compare whole +// combinations. +func assumptionKey(p entity.SpeculationPath) string { + var b strings.Builder + for _, dep := range p.Dependencies { + b.WriteString(dep.Batch) + b.WriteByte('=') + b.WriteString(string(dep.Assumption)) + b.WriteByte(';') + } + return b.String() +} + +// pathIDs renders each candidate's path ID, in order. +func pathIDs(cands []entity.CandidatePath) []string { + out := make([]string, 0, len(cands)) + for _, c := range cands { + out = append(out, c.Path.ID()) + } + return out +} + +// heads renders each candidate's head, in order. +func heads(cands []entity.CandidatePath) []string { + out := make([]string, 0, len(cands)) + for _, c := range cands { + out = append(out, c.Path.Head) + } + return out +} + +// iteratorOf reaches into the concrete iterator, so laziness invariants can be +// asserted on how many paths have actually been built. +func iteratorOf(t *testing.T, iter generator.PathIterator) *iterator { + t.Helper() + it, ok := iter.(*iterator) + require.True(t, ok, "iterator is not the bestfirst iterator") + return it +} + +// countingScorer records how many times each batch is scored. +type countingScorer struct { + scores map[string]float64 + calls map[string]int + total int +} + +func newCountingScorer(scores map[string]float64) *countingScorer { + return &countingScorer{scores: scores, calls: map[string]int{}} +} + +func (c *countingScorer) Score(_ context.Context, b entity.Batch) (float64, error) { + c.calls[b.ID]++ + c.total++ + if v, ok := c.scores[b.ID]; ok { + return v, nil + } + return 0.5, nil +} + +// errScorer always fails, to exercise error propagation from scoring. +type errScorer struct{} + +func (errScorer) Score(context.Context, entity.Batch) (float64, error) { + return 0, assert.AnError +} + +// constScorer scores every batch identically, regardless of ID. +type constScorer struct{ v float64 } + +func (c constScorer) Score(context.Context, entity.Batch) (float64, error) { return c.v, nil } + +// wideHead builds one Speculating head over n unresolved dependencies, each at a +// distinct score so no two combinations tie. +func wideHead(n int) ([]entity.Batch, scorer.Scorer) { + head := entity.Batch{ID: "q/head", State: entity.BatchStateSpeculating} + batches := []entity.Batch{} + scores := map[string]float64{} + for i := 0; i < n; i++ { + dep := fmt.Sprintf("q/d%02d", i) + head.Dependencies = append(head.Dependencies, dep) + // Created deps are unresolved (so they're scored) but not eligible heads. + batches = append(batches, entity.Batch{ID: dep, State: entity.BatchStateCreated}) + scores[dep] = 0.55 + 0.02*float64(i) + } + return append(batches, head), scored(scores) +} + +func TestBestFirst_OrderingAndEnumeration(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating}, + {ID: "q/3", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1", "q/2"}}, + } + sc := scored(map[string]float64{"q/1": 0.9, "q/2": 0.8}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/3") + + // Two unresolved dependencies -> 2^2 candidates. + require.Len(t, cands, 4) + + // Best-first: the all-succeeds path leads, scoring the product of the + // dependency scores (0.9 * 0.8), and scores descend from there. + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(cands[0].Path, "q/1")) + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(cands[0].Path, "q/2")) + assert.InDelta(t, math.Log(0.72), cands[0].RankingScore, 1e-9) + for i := 1; i < len(cands); i++ { + assert.LessOrEqual(t, cands[i].RankingScore, cands[i-1].RankingScore) + } + // The least optimistic candidate excludes both dependencies: (1-0.9)(1-0.8). + last := cands[len(cands)-1] + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(last.Path, "q/1")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(last.Path, "q/2")) + assert.InDelta(t, math.Log(0.02), last.RankingScore, 1e-9) +} + +func TestBestFirst_PinsResolvedDependencies(t *testing.T) { + // A resolved dependency is a fact: its assumption is forced by its + // state and it never varies across the head's paths. + tests := []struct { + name string + state entity.BatchState + want entity.DependencyAssumption + }{ + {name: "succeeded dependency forced to succeeds", state: entity.BatchStateSucceeded, want: entity.DependencyAssumptionSucceeds}, + {name: "failed dependency forced to fails", state: entity.BatchStateFailed, want: entity.DependencyAssumptionFails}, + {name: "cancelled dependency forced to fails", state: entity.BatchStateCancelled, want: entity.DependencyAssumptionFails}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: tt.state}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1"}}, + } + iter, err := New(scored(nil)).Open(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/2") + + // The pinned dependency has no alternative, so there is exactly one + // path, and it carries the forced assumption. + require.Len(t, cands, 1) + assert.Equal(t, tt.want, assumptionFor(cands[0].Path, "q/1")) + }) + } +} + +func TestBestFirst_ResolvedDependenciesDropOutOfSearch(t *testing.T) { + // Three dependencies, but two are resolved facts: only q/open is still + // unresolved, so the head yields 2 paths rather than 8. + batches := []entity.Batch{ + {ID: "q/succeeded", State: entity.BatchStateSucceeded}, + {ID: "q/failed", State: entity.BatchStateFailed}, + {ID: "q/open", State: entity.BatchStateCreated}, + {ID: "q/H", State: entity.BatchStateSpeculating, + Dependencies: []string{"q/succeeded", "q/failed", "q/open"}}, + } + iter, err := New(scored(map[string]float64{"q/open": 0.7})). + Open(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/H") + + require.Len(t, cands, 2) + // Pinned assumptions are identical on both paths and stay in queue order; the + // score reflects only the open dependency, since a fact is a certainty. + assert.Equal(t, "q/succeeded=succeeds;q/failed=fails;q/open=succeeds;", assumptionKey(cands[0].Path)) + assert.InDelta(t, math.Log(0.7), cands[0].RankingScore, 1e-9) + assert.Equal(t, "q/succeeded=succeeds;q/failed=fails;q/open=fails;", assumptionKey(cands[1].Path)) + assert.InDelta(t, math.Log(0.3), cands[1].RankingScore, 1e-9) +} + +func TestBestFirst_EmitsExactSequenceAcrossHeads(t *testing.T) { + // q/solo has nothing to wait on; q/pair assumes outcomes for two in-flight batches. The + // two heads interleave by score rather than draining one at a time. + batches := []entity.Batch{ + {ID: "q/d0", State: entity.BatchStateCreated}, + {ID: "q/d1", State: entity.BatchStateCreated}, + {ID: "q/solo", State: entity.BatchStateSpeculating}, + {ID: "q/pair", State: entity.BatchStateSpeculating, Dependencies: []string{"q/d0", "q/d1"}}, + } + sc := scored(map[string]float64{"q/d0": 0.9, "q/d1": 0.8}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + want := []struct { + head string + assumptions string + probability float64 + }{ + {head: "q/solo", assumptions: "", probability: 1.0}, + {head: "q/pair", assumptions: "q/d0=succeeds;q/d1=succeeds;", probability: 0.72}, + {head: "q/pair", assumptions: "q/d0=succeeds;q/d1=fails;", probability: 0.18}, + {head: "q/pair", assumptions: "q/d0=fails;q/d1=succeeds;", probability: 0.08}, + {head: "q/pair", assumptions: "q/d0=fails;q/d1=fails;", probability: 0.02}, + } + require.Len(t, cands, len(want)) + for i, w := range want { + assert.Equal(t, w.head, cands[i].Path.Head, "head at %d", i) + assert.Equal(t, w.assumptions, assumptionKey(cands[i].Path), "assumptions at %d", i) + assert.InDelta(t, math.Log(w.probability), cands[i].RankingScore, 1e-9, "score at %d", i) + } +} + +func TestBestFirst_PreferredAssumptionFollowsScore(t *testing.T) { + // q/low is more likely to fail than to succeed, so the head's best path assumes + // it fails. The leading path is the preferred-assumption path, not the + // all-succeeds one. + batches := []entity.Batch{ + {ID: "q/high", State: entity.BatchStateCreated}, + {ID: "q/low", State: entity.BatchStateCreated}, + {ID: "q/H", State: entity.BatchStateSpeculating, + Dependencies: []string{"q/high", "q/low"}}, + } + sc := scored(map[string]float64{"q/high": 0.8, "q/low": 0.3}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + want := []struct { + assumptions string + probability float64 + }{ + {assumptions: "q/high=succeeds;q/low=fails;", probability: 0.56}, + {assumptions: "q/high=succeeds;q/low=succeeds;", probability: 0.24}, + {assumptions: "q/high=fails;q/low=fails;", probability: 0.14}, + {assumptions: "q/high=fails;q/low=succeeds;", probability: 0.06}, + } + require.Len(t, cands, len(want)) + for i, w := range want { + assert.Equal(t, w.assumptions, assumptionKey(cands[i].Path), "assumptions at %d", i) + assert.InDelta(t, math.Log(w.probability), cands[i].RankingScore, 1e-9, "score at %d", i) + } + + // Spelled out: the all-succeeds path exists but only ranks second. + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(cands[0].Path, "q/low")) +} + +func TestBestFirst_OnlySpeculatingHeadsProduceCandidates(t *testing.T) { + // Batches in any state supply facts about their dependents, but only a + // Speculating batch is a head worth proposing work on. + tests := []struct { + state entity.BatchState + want bool + }{ + {state: entity.BatchStateUnknown}, + {state: entity.BatchStateCreated}, + {state: entity.BatchStateSpeculating, want: true}, + {state: entity.BatchStateMerging}, + {state: entity.BatchStateSucceeded}, + {state: entity.BatchStateFailed}, + {state: entity.BatchStateCancelling}, + {state: entity.BatchStateCancelled}, + } + + for _, tt := range tests { + name := string(tt.state) + if name == "" { + name = "unknown" + } + t.Run(name, func(t *testing.T) { + batches := []entity.Batch{{ID: "q/1", State: tt.state}} + + iter, err := New(scored(nil)).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + if !tt.want { + assert.Empty(t, cands) + return + } + require.Len(t, cands, 1) + assert.Equal(t, "q/1", cands[0].Path.Head) + }) + } +} + +func TestBestFirst_HeadWithNoDependencies(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + } + + iter, err := New(scored(nil)).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + require.Len(t, cands, 1) + assert.Equal(t, "q/1", cands[0].Path.Head) + assert.Empty(t, cands[0].Path.Dependencies) + assert.InDelta(t, math.Log(1.0), cands[0].RankingScore, 1e-9) +} + +func TestBestFirst_PropagatesScorerError(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1"}}, + } + + iter, err := New(errScorer{}).Open(context.Background(), batches) + assert.Error(t, err) + assert.Nil(t, iter) +} + +func TestBestFirst_AbsentDependencyUsesDefaultScore(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/missing"}}, + } + + // q/missing is not among the live batches, so it cannot be scored at all — + // the generator uses the default score without consulting the + // scorer, whatever the scorer would have said. + iter, err := New(constScorer{0.1}).Open(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/H") + + // Assumed almost certain to succeed, so the best path assumes it does. + require.Len(t, cands, 2) + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(cands[0].Path, "q/missing")) + assert.InDelta(t, math.Log(defaultProbability), cands[0].RankingScore, 1e-9) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(cands[1].Path, "q/missing")) + assert.InDelta(t, math.Log(1-defaultProbability), cands[1].RankingScore, 1e-9) +} + +func TestBestFirst_GeneratesOnlyWhatIsPulled(t *testing.T) { + // 12 unresolved dependencies is an outcome space of 4096 paths. + const deps, space = 12, 1 << 12 + batches, sc := wideHead(deps) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + it := iteratorOf(t, iter) + + // Everything ever built is either handed out already or still waiting in + // the heap. Open builds the head's best path and nothing more. + require.Equal(t, 1, it.heap.Len()) + + pulled := 0 + for i := 0; i < 3; i++ { + _, ok, err := iter.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + pulled++ + } + + // Each pull builds at most two follow-ups, so three pulls can have built at + // most 1+2*3 paths — a tiny fraction of the space. + built := pulled + it.heap.Len() + assert.LessOrEqual(t, built, 7) + assert.Less(t, built, space) +} + +func TestBestFirst_DrainYieldsEveryCombinationOnce(t *testing.T) { + deps := []string{"q/d0", "q/d1", "q/d2"} + batches := []entity.Batch{ + {ID: "q/d0", State: entity.BatchStateCreated}, + {ID: "q/d1", State: entity.BatchStateCreated}, + {ID: "q/d2", State: entity.BatchStateCreated}, + {ID: "q/head", State: entity.BatchStateSpeculating, Dependencies: deps}, + } + sc := scored(map[string]float64{"q/d0": 0.9, "q/d1": 0.7, "q/d2": 0.6}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + require.Len(t, cands, 8) + + // Every include/exclude combination appears exactly once, and each path + // carries its assumptions in queue order. + seen := map[string]int{} + for _, c := range cands { + got := make([]string, 0, len(c.Path.Dependencies)) + for _, dep := range c.Path.Dependencies { + got = append(got, dep.Batch) + } + assert.Equal(t, deps, got, "assumptions must stay in dependency order") + seen[assumptionKey(c.Path)]++ + } + for mask := 0; mask < 8; mask++ { + var want strings.Builder + for i, dep := range deps { + assumption := entity.DependencyAssumptionFails + if mask&(1< 0 { + assert.LessOrEqual(t, cands[i].RankingScore, cands[i-1].RankingScore) + } + } +} + +func TestBestFirst_ResolvedDependenciesAreNeverScored(t *testing.T) { + // A finished build has no probability left to estimate, so the scorer is + // never asked about one. Only the dependency still in flight is scored, and + // only it opens a second path. + batches := []entity.Batch{ + {ID: "q/passed", State: entity.BatchStateSucceeded}, + {ID: "q/broke", State: entity.BatchStateFailed}, + {ID: "q/stopped", State: entity.BatchStateCancelled}, + {ID: "q/running", State: entity.BatchStateCreated}, + {ID: "q/head", State: entity.BatchStateSpeculating, + Dependencies: []string{"q/passed", "q/broke", "q/stopped", "q/running"}}, + } + sc := newCountingScorer(map[string]float64{"q/running": 0.8}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + assert.Equal(t, 1, sc.total, "only the unresolved dependency is scored") + assert.Equal(t, 1, sc.calls["q/running"]) + require.Len(t, cands, 2, "one unresolved dependency, so two paths") + + // The resolved three keep their forced assumption on both paths. + for _, c := range cands { + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(c.Path, "q/passed")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(c.Path, "q/broke")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(c.Path, "q/stopped")) + } +} + +func TestBestFirst_UnknownDependencyStaysUnresolved(t *testing.T) { + // A dependency that is not a live batch has no state to read and no score + // to fetch. It must not be mistaken for resolved: both outcomes stay in the + // space, held apart by the default score. + batches := []entity.Batch{ + {ID: "q/head", State: entity.BatchStateSpeculating, Dependencies: []string{"q/ghost"}}, + } + sc := newCountingScorer(map[string]float64{}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + assert.Zero(t, sc.total, "an unknown dependency is never handed to the scorer") + require.Len(t, cands, 2, "both outcomes stay in the space") + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(cands[0].Path, "q/ghost")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(cands[1].Path, "q/ghost")) + assert.InDelta(t, math.Log(defaultProbability), cands[0].RankingScore, 1e-12) +} + +func TestBestFirst_ReturnedPathsAreIndependent(t *testing.T) { + // Paths are built one at a time from state the head shares, so a caller + // scribbling on what it was handed must not reach the paths still to come. + batches, _ := wideHead(3) + iter, err := New(scored(map[string]float64{"q/d0": 0.9, "q/d1": 0.8, "q/d2": 0.7})). + Open(context.Background(), batches) + require.NoError(t, err) + + first, ok, err := iter.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + before := assumptionKey(first.Path) + + for i := range first.Path.Dependencies { + first.Path.Dependencies[i].Batch = "clobbered" + first.Path.Dependencies[i].Assumption = entity.DependencyAssumptionIgnored + } + + rest := drainAll(t, iter) + require.NotEmpty(t, rest) + for _, c := range rest { + assert.NotContains(t, assumptionKey(c.Path), "clobbered") + assert.NotContains(t, assumptionKey(c.Path), string(entity.DependencyAssumptionIgnored)) + } + assert.NotEqual(t, before, assumptionKey(first.Path), "the test mutated what it was handed") +} + +func TestBestFirst_ScoresAreSummedFromTheHeadsBestPath(t *testing.T) { + // Every score is summed from the head's best-path score in ascending + // alternative order. Deriving one from the path it grew out of instead — + // subtracting the alternative being moved and adding its replacement — + // gives a different float, because floating-point addition does not + // associate, and the heap's ordering guarantee rests on the two being + // summed the same way. The expected scores below come from scoreFor, so a + // walk that stops going through it is what this catches. + deps := []string{"q/d0", "q/d1", "q/d2", "q/d3"} + scores := map[string]float64{"q/d0": 0.9, "q/d1": 0.7, "q/d2": 0.61, "q/d3": 0.5003} + + batches := []entity.Batch{{ + ID: "q/head", State: entity.BatchStateSpeculating, Dependencies: deps, + }} + byID := map[string]entity.Batch{"q/head": batches[0]} + for _, d := range deps { + b := entity.Batch{ID: d, State: entity.BatchStateCreated} + batches = append(batches, b) + byID[d] = b + } + + // Build the same head separately and total every subset the canonical way. + h, err := newHead(context.Background(), byID["q/head"], newDependencyCache(scored(scores), byID)) + require.NoError(t, err) + h.resolveAlternatives() + + want := map[float64]int{} + for mask := 0; mask < 1< b: + return 1 + default: + return 0 + } +} + +// scoreFor returns the score of the path taking the given alternatives. +// takenAlternatives holds positions in the head's cheapest-first alternatives: +// +// [] bestPathScore +// [0] bestPathScore + alternatives[0].relativeScore +// [0 2] bestPathScore + alternatives[0].relativeScore + alternatives[2].relativeScore +// +// Always summed in that order from bestPathScore, never carried over from the +// parent node. Reaching [0 2] from [0 1] by adjusting the parent's score — +// subtracting alternatives[1] and adding alternatives[2] — gives a different +// float from the sum above, because floating-point addition does not +// associate, and expand's ordering guarantee needs both summed the same way. +func (h *head) scoreFor(takenAlternatives []int) float64 { + score := h.bestPathScore + for _, i := range takenAlternatives { + score += h.alternatives[i].relativeScore + } + return score +} + +// pathFor builds the path taking the given alternatives: the head, plus one +// assumption per dependency in queue order — each dependency's preferred +// assumption, or the alternative's where one is taken. +// +// The path owns its dependencies and nothing about the head changes as paths +// are built, so a caller may do as it likes with what it is handed. +// +// Only a head whose alternatives are resolved can have a non-empty set of them +// taken, so the alternatives indexed here are always there to read. +func (h *head) pathFor(takenAlternatives []int) entity.SpeculationPath { + dependencies := make([]entity.PathDependency, len(h.dependencyIDs)) + for i, dep := range h.dependencyIDs { + d, _ := h.cache.known(dep) + dependencies[i] = entity.PathDependency{Batch: dep, Assumption: d.preferred} + } + for _, i := range takenAlternatives { + a := h.alternatives[i] + dependencies[a.depIndex].Assumption = a.assumption + } + return entity.SpeculationPath{Head: h.id, Dependencies: dependencies} +} diff --git a/submitqueue/extension/speculation/generator/bestfirst/iterator.go b/submitqueue/extension/speculation/generator/bestfirst/iterator.go new file mode 100644 index 00000000..b4133cb2 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/iterator.go @@ -0,0 +1,195 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package bestfirst + +import ( + "container/heap" + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// node is one entry in the iterator's heap, standing for a single path in its +// head's speculation space. It names that path by which of the head's +// alternatives it takes rather than carrying a built one, so a node is a +// pointer, a few small ints and a float however many dependencies the head has. +// The path itself is built only if the node is handed out. +type node struct { + // head is the head whose tree this node sits in, and whose alternatives + // takenAlternatives indexes. + head *head + // takenAlternatives holds ascending positions in head.alternatives: the + // dependencies this path assumes the unpreferred way. Empty at the root. + takenAlternatives []int + // score is the path's score, from head.scoreFor. Held rather than + // recomputed on every heap comparison. + score float64 +} + +// depth is how far the node has deviated from its head's best path. +func (n node) depth() int { return len(n.takenAlternatives) } + +// nodeHeap keeps nodes best-first: highest score first. Scores are log +// probabilities, so "highest" means closest to 0 and negative infinity sorts +// last. Go's container/heap does the real work here; the methods below are +// just the interface it asks for. +// +// Equal scores break on depth first, then on the alternatives taken, then on +// head. Depth first offers every head's root before any head's deeper nodes, +// which matters because a dependency at probability exactly 0.5 deviates for +// free: break on head instead and one head's whole coin-flip subtree drains +// the caller's budget before another head's root is ever seen. Comparing the +// alternatives before the head then interleaves heads at equal depth, since +// position i means "this head's i-th cheapest deviation" either way. +// +// The tie-break makes a run repeatable, not globally ordered: only nodes in the +// heap at the same moment are ever compared. +type nodeHeap []node + +func (h nodeHeap) Len() int { return len(h) } +func (h nodeHeap) Less(i, j int) bool { + a, b := h[i], h[j] + // Greater-than, not less-than: that is what makes this a max-heap, so the + // highest-scoring path ends up on top. + if a.score != b.score { + return a.score > b.score + } + if a.depth() != b.depth() { + return a.depth() < b.depth() + } + // Same depth, so the two sets are the same length and this reads both. + for n := range a.takenAlternatives { + if a.takenAlternatives[n] != b.takenAlternatives[n] { + return a.takenAlternatives[n] < b.takenAlternatives[n] + } + } + return a.head.id < b.head.id +} +func (h nodeHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *nodeHeap) Push(x any) { *h = append(*h, x.(node)) } +func (h *nodeHeap) Pop() any { + // Careful: this does not pick the best entry. container/heap has already + // swapped the best one to the end of the slice before calling us, so all we + // do is chop the last element off and hand it back. + old := *h + n := len(old) + e := old[n-1] + *h = old[:n-1] + return e +} + +// iterator hands out the queue's paths, best first, building each as it goes. +// +// Every head roots a tree of its paths, and one heap holds the frontier of all +// of them at once: the nodes reached but not yet handed out. Handing one out +// adds its children, so the heap always holds each head's current best, and +// taking the top of it is taking the best path in the queue. The README works +// through a full example. +type iterator struct { + heap nodeHeap +} + +// Next hands out the best remaining path across all heads. On the way out it +// adds that node's children, so the heap always has the next ones ready, and +// builds the one path it returns. It returns ok=false once nothing is left. +// +// Nothing after the pop can fail, so a node is never taken off the heap and +// then dropped. +// +// A cancelled or expired ctx ends the stream with its error and no candidate. +func (it *iterator) Next(ctx context.Context) (entity.CandidatePath, bool, error) { + if err := ctx.Err(); err != nil { + return entity.CandidatePath{}, false, err + } + if it.heap.Len() == 0 { + return entity.CandidatePath{}, false, nil + } + n := heap.Pop(&it.heap).(node) + it.expand(n) + return entity.CandidatePath{ + Path: n.head.pathFor(n.takenAlternatives), + RankingScore: n.score, + }, true, nil +} + +// expand adds a node's children — the paths that come just after it within the +// same head. +// +// A node takes an ascending set of positions in its head's cheapest-first +// alternatives; the root takes none, and a head with n alternatives has 2^n +// nodes in its tree, one per path. Treat the set as something that only moves +// forward: if the highest position taken so far is i, there are two ways on: +// +// take i+1 too {0,1} -> {0,1,2} one more dependency assumed the other way +// swap i for i+1 {0,1} -> {0,2} move the last one along +// +// Every node is reached exactly once. You can always tell which move produced a +// set — its last two positions are adjacent iff it was the first move — so each +// node has exactly one parent, and re-tracing the moves from the root reaches +// every set. That is why nothing needs to remember where it has been. +// +// A child never outscores its parent. Both moves add a relativeScore no greater +// than the one they replace or build on, since alternatives are sorted cheapest +// first and every relativeScore is at most 0, and scoreFor sums both from +// bestPathScore over the identical prefix — so the ordering holds in floating +// point as computed, not just in exact arithmetic. Carrying a score over from +// the parent instead would break that: advancing {0,1} to {0,2} would give +// (s-r1)+r2 rather than bestPathScore+r0+r2, and floating-point addition does +// not associate. +// +// Since a node is the heap's maximum when it is popped and its children never +// outscore it, the paths come out in non-increasing score order. +func (it *iterator) expand(n node) { + // A head's alternatives are worked out when its root is expanded, which is + // this call if n is that root. + n.head.resolveAlternatives() + + taken := n.takenAlternatives + k := n.depth() + last := -1 + if k > 0 { + last = taken[k-1] + } + // Every alternative past the last taken one is already used up, so there is + // no move left to make and this node is a leaf. + if last+1 >= len(n.head.alternatives) { + return + } + + // Take one more alternative, keeping the ones already taken. + extended := make([]int, k+1) + copy(extended, taken) + extended[k] = last + 1 + it.push(n.head, extended) + + // Move the last taken alternative one position along. A root has none to + // move, so this second child doesn't apply to it. + if k > 0 { + advanced := make([]int, k) + copy(advanced, taken) + advanced[k-1] = last + 1 + it.push(n.head, advanced) + } +} + +// push adds the node for the given alternatives to the heap. It scores the node +// but does not build its path — that waits until the path is handed out. +func (it *iterator) push(h *head, takenAlternatives []int) { + heap.Push(&it.heap, node{ + head: h, + takenAlternatives: takenAlternatives, + score: h.scoreFor(takenAlternatives), + }) +} diff --git a/submitqueue/extension/speculation/generator/generator.go b/submitqueue/extension/speculation/generator/generator.go new file mode 100644 index 00000000..9243fff4 --- /dev/null +++ b/submitqueue/extension/speculation/generator/generator.go @@ -0,0 +1,57 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package generator defines the candidate-stream composition point used by the +// standard Speculator. A Generator yields candidate paths through a pull +// iterator, in an order it chooses; it is not a controller-facing extension — +// an alternate Speculator need not use or expose this split. +package generator + +//go:generate mockgen -source=generator.go -destination=mock/generator_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// Generator produces the queue's candidate paths as a pull-based stream over the +// queue's live batches. The order candidates are yielded in, and how they are +// ranked, is the generator's own policy. +type Generator interface { + // Open starts the candidate stream. The consumer pulls candidates lazily from + // the returned iterator. A cancelled or expired ctx aborts with its error + // and no iterator. + // + // batches is treated as a snapshot: a generator may hold it, and the + // dependency slices within it, for as long as the iterator lives rather + // than copying them. Callers must not write to either while pulling. A + // queue that has moved on is a new snapshot and a new Open, which is how + // batches are revised anyway — they are replaced, not edited in place. + Open(ctx context.Context, batches []entity.Batch) (PathIterator, error) +} + +// PathIterator is a pull-based stream of candidate paths. Beyond what ranking +// requires up front, the producer computes only what is pulled. +type PathIterator interface { + // Next yields the next candidate in the generator's order. ok is false once + // the generator has no more candidates to offer. Candidates never repeat and + // never contradict a resolved fact — a dependency that already succeeded is + // never assumed to fail, and one that already failed is never assumed to + // succeed. + // + // A cancelled or expired ctx ends the stream with its error, ok false, and + // no candidate. + Next(ctx context.Context) (c entity.CandidatePath, ok bool, err error) +} diff --git a/submitqueue/extension/speculation/generator/mock/BUILD.bazel b/submitqueue/extension/speculation/generator/mock/BUILD.bazel new file mode 100644 index 00000000..9305c151 --- /dev/null +++ b/submitqueue/extension/speculation/generator/mock/BUILD.bazel @@ -0,0 +1,13 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["generator_mock.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/generator/mock", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + "@org_uber_go_mock//gomock:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/generator/mock/generator_mock.go b/submitqueue/extension/speculation/generator/mock/generator_mock.go new file mode 100644 index 00000000..6b1dea8c --- /dev/null +++ b/submitqueue/extension/speculation/generator/mock/generator_mock.go @@ -0,0 +1,98 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: generator.go +// +// Generated by this command: +// +// mockgen -source=generator.go -destination=mock/generator_mock.go -package=mock +// + +// Package mock is a generated GoMock package. +package mock + +import ( + context "context" + reflect "reflect" + + entity "github.com/uber/submitqueue/submitqueue/entity" + generator "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" + gomock "go.uber.org/mock/gomock" +) + +// MockGenerator is a mock of Generator interface. +type MockGenerator struct { + ctrl *gomock.Controller + recorder *MockGeneratorMockRecorder + isgomock struct{} +} + +// MockGeneratorMockRecorder is the mock recorder for MockGenerator. +type MockGeneratorMockRecorder struct { + mock *MockGenerator +} + +// NewMockGenerator creates a new mock instance. +func NewMockGenerator(ctrl *gomock.Controller) *MockGenerator { + mock := &MockGenerator{ctrl: ctrl} + mock.recorder = &MockGeneratorMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockGenerator) EXPECT() *MockGeneratorMockRecorder { + return m.recorder +} + +// Open mocks base method. +func (m *MockGenerator) Open(ctx context.Context, batches []entity.Batch) (generator.PathIterator, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Open", ctx, batches) + ret0, _ := ret[0].(generator.PathIterator) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Open indicates an expected call of Open. +func (mr *MockGeneratorMockRecorder) Open(ctx, batches any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Open", reflect.TypeOf((*MockGenerator)(nil).Open), ctx, batches) +} + +// MockPathIterator is a mock of PathIterator interface. +type MockPathIterator struct { + ctrl *gomock.Controller + recorder *MockPathIteratorMockRecorder + isgomock struct{} +} + +// MockPathIteratorMockRecorder is the mock recorder for MockPathIterator. +type MockPathIteratorMockRecorder struct { + mock *MockPathIterator +} + +// NewMockPathIterator creates a new mock instance. +func NewMockPathIterator(ctrl *gomock.Controller) *MockPathIterator { + mock := &MockPathIterator{ctrl: ctrl} + mock.recorder = &MockPathIteratorMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockPathIterator) EXPECT() *MockPathIteratorMockRecorder { + return m.recorder +} + +// Next mocks base method. +func (m *MockPathIterator) Next(ctx context.Context) (entity.CandidatePath, bool, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Next", ctx) + ret0, _ := ret[0].(entity.CandidatePath) + ret1, _ := ret[1].(bool) + ret2, _ := ret[2].(error) + return ret0, ret1, ret2 +} + +// Next indicates an expected call of Next. +func (mr *MockPathIteratorMockRecorder) Next(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Next", reflect.TypeOf((*MockPathIterator)(nil).Next), ctx) +}