diff --git a/submitqueue/entity/speculation.go b/submitqueue/entity/speculation.go index bb4cb4f4..1794d5b8 100644 --- a/submitqueue/entity/speculation.go +++ b/submitqueue/entity/speculation.go @@ -196,8 +196,8 @@ type CandidatePath struct { // Path is the candidate: a head plus one assumption per dependency. Path SpeculationPath // RankingScore is the score the Generator ranked this candidate by. Higher - // sorts first. The scale is the Generator's own — consumers order by it but - // do not interpret it — and it is meaningful only within the run that - // produced it, which is why it is never stored. + // sorts first within the Generator's own ranking. Consumers take candidates + // in iterator order and do not interpret the value. It is meaningful only + // within the run that produced it, which is why it is never stored. RankingScore float64 } 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..61a0613a --- /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 ordered stream 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. + +`Generate` starts the stream over the queue's live batches and returns an `Iterator`. Beyond any ranking work required up front, the generator computes only the candidates the caller pulls. A cancelled or expired context ends the stream with its error. The snapshot must include every batch a head's direct dependencies reference; a snapshot that does not — or that carries empty or duplicate IDs, or a head with an empty, duplicate, or self dependency — is malformed input and errors instead of opening a stream. + +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..ff2cdd21 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel @@ -0,0 +1,26 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["bestfirst.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..f43e83cd --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/README.md @@ -0,0 +1,15 @@ +# bestfirst + +`bestfirst` implements `generator.Generator` by ranking candidate paths by the probability that all their dependency assumptions hold. It returns one path per pull across all speculating heads without enumerating every combination up front. + +The [best-first speculation path generation RFC](../../../../../doc/rfc/submitqueue/speculation-generator-best-first.md) defines the terminology, algorithm, correctness argument, worked example, and alternatives considered. This README records only the package's operational behavior. + +## Behavior + +- `Generate` validates the snapshot, scores each unique unresolved direct dependency once, fixes assumptions for resolved dependencies, calculates each head's best score, and seeds the global heap with every eligible head's best-path candidate. Each head's remaining paths wait in that head's own lazy stream, whose flips are worked out only when the head is first handed out. +- `Next` removes the highest-ranked candidate, advances only that head's stream, constructs that candidate's complete path, and returns it. Pulling long enough returns every path exactly once in non-increasing score order. +- Ranking scores are sums of log probabilities, avoiding underflow while preserving probability order. They are meaningful only within the run that produced them. +- Exact ties prefer fewer flips, then compare taken flip indexes, then head ID. +- The snapshot must contain every batch a head's direct dependencies reference; a snapshot missing one — or carrying empty or duplicate batch IDs, or a head with an empty, duplicate, or self dependency — is malformed input and errors instead of opening a stream. Any defaulting for a batch that is hard to score belongs to the scorer, not the generator. + +The behavior is covered by `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..9454e305 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go @@ -0,0 +1,442 @@ +// 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 provides a probability-ordered speculation path generator. +// +// Throughout, a probability is the [0, 1] value a scorer gives; a score is its +// logarithm. Scores are summed and compared, never exponentiated, so wide +// heads cannot underflow into ties. The algorithm — per-head streams +// enumerating flip subsets lazily, merged through one global heap — is +// documented in doc/rfc/submitqueue/speculation-generator-best-first.md. +package bestfirst + +import ( + "container/heap" + "context" + "fmt" + "math" + "slices" + "sort" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/scorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +// bestFirst generates candidate paths using independent dependency +// probabilities supplied by scorer. +type bestFirst struct { + scorer scorer.Scorer +} + +var _ generator.Generator = (*bestFirst)(nil) + +// New returns a Generator that ranks paths by the probability that every +// unresolved dependency assumption holds. The scorer is called at most once +// per unresolved dependency batch in each Generate call. +func New(s scorer.Scorer) generator.Generator { + if s == nil { + panic("bestfirst.New: scorer must not be nil") + } + return &bestFirst{scorer: s} +} + +// Generate validates the queue snapshot, scores the unresolved dependencies of +// Speculating heads, and opens a lazy global best-first iterator. +func (g *bestFirst) Generate(ctx context.Context, batches []entity.Batch) (generator.Iterator, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + batchByID := make(map[string]entity.Batch, len(batches)) + for _, batch := range batches { + if batch.ID == "" { + return nil, fmt.Errorf("batch has an empty ID") + } + if _, exists := batchByID[batch.ID]; exists { + return nil, fmt.Errorf("duplicate batch ID %q", batch.ID) + } + batchByID[batch.ID] = batch + } + + heads := make([]entity.Batch, 0) + unresolvedDependencyIDs := make(map[string]struct{}) + for _, batch := range batches { + if batch.State != entity.BatchStateSpeculating { + continue + } + heads = append(heads, batch) + + seen := make(map[string]struct{}, len(batch.Dependencies)) + for _, dependencyID := range batch.Dependencies { + if dependencyID == "" { + return nil, fmt.Errorf("head %q has an empty dependency ID", batch.ID) + } + if dependencyID == batch.ID { + return nil, fmt.Errorf("head %q depends on itself", batch.ID) + } + if _, duplicate := seen[dependencyID]; duplicate { + return nil, fmt.Errorf("head %q repeats dependency %q", batch.ID, dependencyID) + } + seen[dependencyID] = struct{}{} + + dependency, exists := batchByID[dependencyID] + if !exists { + return nil, fmt.Errorf("head %q references dependency %q missing from the snapshot", batch.ID, dependencyID) + } + if dependency.State == entity.BatchStateUnknown { + return nil, fmt.Errorf("dependency %q has an unknown state", dependencyID) + } + if _, resolved := resolvedAssumption(dependency.State); !resolved { + unresolvedDependencyIDs[dependencyID] = struct{}{} + } + } + } + sort.Slice(heads, func(i, j int) bool { + return heads[i].ID < heads[j].ID + }) + + // Score each unique unresolved dependency once, in a stable order. A score + // outside [0, 1] is rejected here because everything downstream treats it + // as a probability, and a bad value would corrupt the ordering silently. + ids := make([]string, 0, len(unresolvedDependencyIDs)) + for id := range unresolvedDependencyIDs { + ids = append(ids, id) + } + sort.Strings(ids) + probabilityByID := make(map[string]float64, len(ids)) + for _, id := range ids { + if err := ctx.Err(); err != nil { + return nil, err + } + probability, err := g.scorer.Score(ctx, batchByID[id]) + if err != nil { + return nil, fmt.Errorf("score dependency %q: %w", id, err) + } + if math.IsNaN(probability) || probability < 0 || probability > 1 { + return nil, fmt.Errorf("scorer returned %v for batch %q: want a probability in [0, 1]", probability, id) + } + probabilityByID[id] = probability + } + + it := &candidateIterator{} + heap.Init(&it.candidates) + for _, head := range heads { + stream := newPathStream(head, batchByID, probabilityByID) + heap.Push(&it.candidates, candidateItem{ + stream: stream, + score: stream.bestScore, + }) + } + return it, nil +} + +// resolvedAssumption converts a terminal dependency outcome into the only +// coherent path assumption. Cancelling remains unresolved because cancellation +// is best-effort and the batch may still succeed. +func resolvedAssumption(state entity.BatchState) (entity.DependencyAssumption, bool) { + switch state { + case entity.BatchStateSucceeded: + return entity.DependencyAssumptionSucceeds, true + case entity.BatchStateFailed, entity.BatchStateCancelled: + return entity.DependencyAssumptionFails, true + default: + return entity.DependencyAssumptionUnknown, false + } +} + +// dependencyVariable is one unresolved dependency and the penalty for flipping +// away from its most likely outcome. +type dependencyVariable struct { + // dependencyIndex is the dependency's position in the head's dependency + // list — the slot in a built path a flip rewrites. + dependencyIndex int + // flipCost is what flipping adds to a path's score: + // log(opposite) - log(preferred). At most 0, and -Inf when the opposite + // side cannot happen. + flipCost float64 +} + +// pathStream lazily enumerates one head's paths in descending probability. Its +// local heap holds the flip subsets reached but not yet handed out; the head's +// current best candidate sits in the iterator's global heap instead. +type pathStream struct { + // head is the head batch's ID, the Head of every path the stream yields. + head string + // base is the most likely path's assumptions in queue order: the preferred + // side where unresolved, the forced side where resolved. build copies it + // and flips the taken slots. + base []entity.PathDependency + // variables are the head's unresolved dependencies, sorted by ascending + // flip penalty once prepared. + variables []dependencyVariable + // bestScore is the no-flip path's score: the preferred log probabilities + // summed. Every other path's score is this plus its taken flip costs. + bestScore float64 + // prepared reports that variables are sorted and the local heap is seeded, + // which is deferred until the head is first pulled. + prepared bool + // subsets is the local heap of flip subsets. + subsets flipHeap +} + +func newPathStream(head entity.Batch, batchByID map[string]entity.Batch, probabilityByID map[string]float64) *pathStream { + stream := &pathStream{ + head: head.ID, + base: make([]entity.PathDependency, len(head.Dependencies)), + } + for i, dependencyID := range head.Dependencies { + if assumption, resolved := resolvedAssumption(batchByID[dependencyID].State); resolved { + // A fact: it contributes probability 1 (log 0) and offers no flip. + stream.base[i] = entity.PathDependency{Batch: dependencyID, Assumption: assumption} + continue + } + + // The preferred side is the more probable one; on an exact tie, + // succeeds. Its probability is at least 0.5, so its log is finite. + probability := probabilityByID[dependencyID] + preferredProbability := math.Max(probability, 1-probability) + oppositeProbability := math.Min(probability, 1-probability) + preferred := entity.DependencyAssumptionFails + if probability >= 0.5 { + preferred = entity.DependencyAssumptionSucceeds + } + stream.base[i] = entity.PathDependency{Batch: dependencyID, Assumption: preferred} + stream.bestScore += math.Log(preferredProbability) + stream.variables = append(stream.variables, dependencyVariable{ + dependencyIndex: i, + flipCost: math.Log(oppositeProbability) - math.Log(preferredProbability), + }) + } + return stream +} + +// prepare sorts the variables cheapest-flip-first and seeds the local heap +// with the single cheapest flip — the no-flip path's lone successor. Deferred +// until the stream first advances, so a head nobody pulls never pays the sort. +func (s *pathStream) prepare() { + if s.prepared { + return + } + sort.Slice(s.variables, func(i, j int) bool { + left, right := s.variables[i], s.variables[j] + if left.flipCost != right.flipCost { + // Descending: costs are at most 0 and closest to 0 is cheapest. + return left.flipCost > right.flipCost + } + return left.dependencyIndex < right.dependencyIndex + }) + s.prepared = true + if len(s.variables) > 0 { + s.push([]int{0}) + } +} + +// next returns the head's next-best flip subset and pushes that subset's +// extend and swap successors, so the local heap always holds the subsets that +// could follow. ok is false once the head is exhausted. +// +// The extend/swap tree reaches every subset exactly once, and because the +// variables are sorted and every flipCost is at most 0, no successor outscores +// its parent — in floating point as computed, since scoreFor sums parent and +// child from bestScore over the identical prefix. +func (s *pathStream) next() (flipped []int, score float64, ok bool) { + s.prepare() + if s.subsets.Len() == 0 { + return nil, 0, false + } + entry := heap.Pop(&s.subsets).(flipEntry) + + if j := entry.flipped[len(entry.flipped)-1]; j+1 < len(s.variables) { + s.push(appendCopy(entry.flipped, j+1)) // extend: also flip j+1 + s.push(replaceLastCopy(entry.flipped, j+1)) // swap: trade flip j for j+1 + } + return entry.flipped, entry.score, true +} + +func (s *pathStream) push(flipped []int) { + heap.Push(&s.subsets, flipEntry{flipped: flipped, score: s.scoreFor(flipped)}) +} + +// scoreFor sums the path's score from bestScore in ascending flip order. It is +// never adjusted incrementally from a parent's score: floating-point addition +// does not associate, and the heap ordering depends on parent and child being +// summed the same way. The fresh sum is also why a -Inf flip needs no special +// case — nothing is ever subtracted. +func (s *pathStream) scoreFor(flipped []int) float64 { + score := s.bestScore + for _, i := range flipped { + score += s.variables[i].flipCost + } + return score +} + +// build constructs the path taking the given flips. The returned path owns its +// dependencies. +func (s *pathStream) build(flipped []int) entity.SpeculationPath { + dependencies := make([]entity.PathDependency, len(s.base)) + copy(dependencies, s.base) + for _, i := range flipped { + at := s.variables[i].dependencyIndex + dependencies[at].Assumption = opposite(dependencies[at].Assumption) + } + return entity.SpeculationPath{Head: s.head, Dependencies: dependencies} +} + +// opposite is the other side of an assumption. Flips only apply to unresolved +// dependencies, whose base assumption is exactly one of the two sides. +func opposite(assumption entity.DependencyAssumption) entity.DependencyAssumption { + if assumption == entity.DependencyAssumptionSucceeds { + return entity.DependencyAssumptionFails + } + return entity.DependencyAssumptionSucceeds +} + +func appendCopy(values []int, value int) []int { + result := make([]int, len(values)+1) + copy(result, values) + result[len(values)] = value + return result +} + +func replaceLastCopy(values []int, value int) []int { + result := make([]int, len(values)) + copy(result, values) + result[len(result)-1] = value + return result +} + +// flipEntry is one subset in a stream's local heap: the flips a path takes, as +// ascending indexes into the sorted variables, plus the path's score. +type flipEntry struct { + flipped []int + score float64 +} + +// flipHeap orders a stream's subsets best-first: highest score first, ties +// preferring fewer flips and then the lexicographically smaller subset. That +// is a strict total order, so iteration is deterministic without insertion +// counters. +type flipHeap []flipEntry + +var _ heap.Interface = (*flipHeap)(nil) + +func (h flipHeap) Len() int { return len(h) } + +func (h flipHeap) Less(i, j int) bool { + left, right := h[i], h[j] + if left.score != right.score { + return left.score > right.score + } + if len(left.flipped) != len(right.flipped) { + return len(left.flipped) < len(right.flipped) + } + return slices.Compare(left.flipped, right.flipped) < 0 +} + +func (h flipHeap) Swap(i, j int) { + h[i], h[j] = h[j], h[i] +} + +func (h *flipHeap) Push(value any) { + *h = append(*h, value.(flipEntry)) +} + +func (h *flipHeap) Pop() any { + old := *h + last := len(old) - 1 + value := old[last] + *h = old[:last] + return value +} + +// candidateIterator performs a k-way merge of the per-head ordered streams. +type candidateIterator struct { + candidates candidateHeap +} + +var _ generator.Iterator = (*candidateIterator)(nil) + +// Next returns the best remaining path across all heads: it pops the global +// best, advances only that head's stream, and reinserts the head's next +// candidate. The path is built only here, when it is handed out, and nothing +// after the pop can fail — a cancelled ctx does not consume a candidate. +func (i *candidateIterator) Next(ctx context.Context) (entity.CandidatePath, bool, error) { + if err := ctx.Err(); err != nil { + return entity.CandidatePath{}, false, err + } + if i.candidates.Len() == 0 { + return entity.CandidatePath{}, false, nil + } + + item := heap.Pop(&i.candidates).(candidateItem) + if flipped, score, ok := item.stream.next(); ok { + heap.Push(&i.candidates, candidateItem{stream: item.stream, flipped: flipped, score: score}) + } + return entity.CandidatePath{ + Path: item.stream.build(item.flipped), + RankingScore: item.score, + }, true, nil +} + +// candidateItem is one head's current best: the stream it advances, the flips +// the path takes, and its score. flipped is nil for the head's no-flip path. +type candidateItem struct { + stream *pathStream + flipped []int + score float64 +} + +// candidateHeap is a max-heap by score, holding one item per live head. +// RankingScore carries the same log value, so ordering never underflows. +// +// Ties prefer fewer flips before comparing heads: a dependency at exactly 0.5 +// flips for free, and breaking on head first would drain one head's whole +// coin-flip subtree before another head's best path was ever offered. The +// remaining comparisons make a run repeatable. +type candidateHeap []candidateItem + +var _ heap.Interface = (*candidateHeap)(nil) + +func (h candidateHeap) Len() int { return len(h) } + +func (h candidateHeap) Less(i, j int) bool { + left, right := h[i], h[j] + if left.score != right.score { + return left.score > right.score + } + if len(left.flipped) != len(right.flipped) { + return len(left.flipped) < len(right.flipped) + } + if c := slices.Compare(left.flipped, right.flipped); c != 0 { + return c < 0 + } + return left.stream.head < right.stream.head +} + +func (h candidateHeap) Swap(i, j int) { + h[i], h[j] = h[j], h[i] +} + +func (h *candidateHeap) Push(value any) { + *h = append(*h, value.(candidateItem)) +} + +func (h *candidateHeap) Pop() any { + old := *h + last := len(old) - 1 + value := old[last] + *h = old[:last] + return value +} 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..d07a27a8 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go @@ -0,0 +1,1081 @@ +// 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.Iterator) []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.Iterator) *candidateIterator { + t.Helper() + it, ok := iter.(*candidateIterator) + 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/dep%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/A", State: entity.BatchStateSpeculating}, + {ID: "q/B", State: entity.BatchStateSpeculating}, + {ID: "q/C", State: entity.BatchStateSpeculating, Dependencies: []string{"q/A", "q/B"}}, + } + sc := scored(map[string]float64{"q/A": 0.9, "q/B": 0.8}) + + iter, err := New(sc).Generate(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/C") + + // 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/A")) + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(cands[0].Path, "q/B")) + 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/A")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(last.Path, "q/B")) + 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/A", State: tt.state}, + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/A"}}, + } + iter, err := New(scored(nil)).Generate(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/H") + + // The pinned dependency has no flip, 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/A")) + }) + } +} + +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})). + Generate(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/D has nothing to wait on; q/C 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/A", State: entity.BatchStateCreated}, + {ID: "q/B", State: entity.BatchStateCreated}, + {ID: "q/D", State: entity.BatchStateSpeculating}, + {ID: "q/C", State: entity.BatchStateSpeculating, Dependencies: []string{"q/A", "q/B"}}, + } + sc := scored(map[string]float64{"q/A": 0.9, "q/B": 0.8}) + + iter, err := New(sc).Generate(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + want := []struct { + head string + assumptions string + probability float64 + }{ + {head: "q/D", assumptions: "", probability: 1.0}, + {head: "q/C", assumptions: "q/A=succeeds;q/B=succeeds;", probability: 0.72}, + {head: "q/C", assumptions: "q/A=succeeds;q/B=fails;", probability: 0.18}, + {head: "q/C", assumptions: "q/A=fails;q/B=succeeds;", probability: 0.08}, + {head: "q/C", assumptions: "q/A=fails;q/B=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 most likely 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).Generate(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/H", State: tt.state}} + + iter, err := New(scored(nil)).Generate(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/H", cands[0].Path.Head) + }) + } +} + +func TestBestFirst_HeadWithNoDependencies(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/H", State: entity.BatchStateSpeculating}, + } + + iter, err := New(scored(nil)).Generate(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + require.Len(t, cands, 1) + assert.Equal(t, "q/H", 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/A", State: entity.BatchStateSpeculating}, + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/A"}}, + } + + iter, err := New(errScorer{}).Generate(context.Background(), batches) + assert.Error(t, err) + assert.Nil(t, iter) +} + +func TestBestFirst_RejectsMalformedSnapshots(t *testing.T) { + // The snapshot contract: every batch a head's direct dependencies reference + // is present with a readable state, IDs are unique and non-empty, and no + // head repeats a dependency or depends on itself. Anything else is + // malformed input. In particular a missing dependency is never guessed + // about — the scorer, not the generator, owns any defaulting for batches + // that are hard to score. + tests := []struct { + name string + batches []entity.Batch + }{ + { + name: "empty batch ID", + batches: []entity.Batch{{State: entity.BatchStateSpeculating}}, + }, + { + name: "duplicate batch ID", + batches: []entity.Batch{ + {ID: "q/A", State: entity.BatchStateCreated}, + {ID: "q/A", State: entity.BatchStateSpeculating}, + }, + }, + { + name: "empty dependency ID", + batches: []entity.Batch{ + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{""}}, + }, + }, + { + name: "self dependency", + batches: []entity.Batch{ + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/H"}}, + }, + }, + { + name: "duplicate dependency", + batches: []entity.Batch{ + {ID: "q/A", State: entity.BatchStateCreated}, + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/A", "q/A"}}, + }, + }, + { + name: "missing dependency", + batches: []entity.Batch{ + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/ghost"}}, + }, + }, + { + name: "unknown dependency state", + batches: []entity.Batch{ + {ID: "q/A", State: entity.BatchStateUnknown}, + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/A"}}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + iter, err := New(scored(nil)).Generate(context.Background(), tt.batches) + require.Error(t, err) + assert.Nil(t, iter) + }) + } +} + +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).Generate(context.Background(), batches) + require.NoError(t, err) + it := iteratorOf(t, iter) + + // Everything ever worked out is either handed out already, waiting in the + // global heap, or waiting in the head's local heap. Generate pushes the + // head's most likely path and nothing more. + require.Equal(t, 1, it.candidates.Len()) + stream := it.candidates[0].stream + require.Empty(t, stream.subsets, "Generate must not touch the local heap") + + pulled := 0 + for i := 0; i < 3; i++ { + _, ok, err := iter.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + pulled++ + } + + // Advancing the stream pushes at most two successor subsets per pull, on + // top of the one-time seed of the cheapest single flip, so three pulls can + // have worked out at most 2+2*3 paths — a tiny fraction of the space. + built := pulled + it.candidates.Len() + len(stream.subsets) + assert.LessOrEqual(t, built, 8) + assert.Less(t, built, space) +} + +func TestBestFirst_DrainYieldsEveryCombinationOnce(t *testing.T) { + deps := []string{"q/A", "q/B", "q/C"} + batches := []entity.Batch{ + {ID: "q/A", State: entity.BatchStateCreated}, + {ID: "q/B", State: entity.BatchStateCreated}, + {ID: "q/C", State: entity.BatchStateCreated}, + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: deps}, + } + sc := scored(map[string]float64{"q/A": 0.9, "q/B": 0.7, "q/C": 0.6}) + + iter, err := New(sc).Generate(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/H", 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).Generate(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_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/dep00": 0.9, "q/dep01": 0.8, "q/dep02": 0.7})). + Generate(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_ScoresAreSummedFromTheHeadsBestScore(t *testing.T) { + // Every score is summed from the head's best score in ascending + // flip order. Deriving one from the path it grew out of instead — + // subtracting the flip 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/A", "q/B", "q/C", "q/D"} + scores := map[string]float64{"q/A": 0.9, "q/B": 0.7, "q/C": 0.61, "q/D": 0.5003} + + batches := []entity.Batch{{ + ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: deps, + }} + byID := map[string]entity.Batch{"q/H": 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. + s := newPathStream(byID["q/H"], byID, scores) + s.prepare() + + want := map[float64]int{} + for mask := 0; mask < 1<