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..6fb8f8f6 --- /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. + +`Open` starts the stream over the queue's live batches and returns a `PathIterator`. 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. + +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..4bf15a8f --- /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 + +- `Open` scores each unique unresolved dependency once, fixes assumptions for resolved dependencies, calculates each head's best score, and seeds one shared heap with every eligible head's best-path node. +- `Next` removes the highest-ranked node, adds at most two child nodes, constructs that node'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 positions, then head ID. +- A missing dependency uses a success probability of 0.95 because it cannot be scored from the input snapshot. + +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..982eb290 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go @@ -0,0 +1,325 @@ +// 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. +// - A flip is the option to take the unpreferred assumption for a +// dependency still unresolved. The head's most likely 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 flips 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 most likely path — every remaining dependency on its preferred +// assumption. That total is the head's opening score. +// 3. Put that most likely path into a single heap shared by every head, as a node: +// the head, plus the (still empty) set of flips 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 +// flips are worked out until something pulls from it. +// +// The RFC in doc/rfc/submitqueue/speculation-generator-best-first.md walks a +// two-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 flips 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 score up front. The RFC gives the resulting cost. +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 switches sides: 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 flipCost 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 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 flips taken, every + // dependency on its preferred assumption. Nothing is built for it yet — + // not the head's flips, 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 flip, 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 more probable + // 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 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 + // flipCost 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 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). + flipCost 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 flip 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 + } + + // The preferred side is the assumption with the higher probability; on an + // exact tie, 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. + flipCost: 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..03c3f5e0 --- /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/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).Open(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)).Open(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})). + 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/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).Open(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).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/H", 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/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)).Open(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{}).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 most likely 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 most likely 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 child nodes, 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/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).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/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).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/H", 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/dep00": 0.9, "q/dep01": 0.8, "q/dep02": 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_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. + h, err := newHead(context.Background(), byID["q/H"], newDependencyCache(scored(scores), byID)) + require.NoError(t, err) + h.resolveFlips() + + 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 flips. +// takenFlips holds positions in the head's cheapest-first flips: +// +// [] bestScore +// [0] bestScore + flips[0].flipCost +// [0 2] bestScore + flips[0].flipCost + flips[2].flipCost +// +// Always summed in that order from bestScore, never carried over from the +// parent node. Reaching [0 2] from [0 1] by adjusting the parent's score — +// subtracting flips[1] and adding flips[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(takenFlips []int) float64 { + score := h.bestScore + for _, i := range takenFlips { + score += h.flips[i].flipCost + } + return score +} + +// pathFor builds the path taking the given flips: the head, plus one +// assumption per dependency in queue order — each dependency's preferred +// assumption, or the flip'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 flips are resolved can have a non-empty set of them +// taken, so the flips indexed here are always there to read. +func (h *head) pathFor(takenFlips []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 takenFlips { + a := h.flips[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..841a0523 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/iterator.go @@ -0,0 +1,196 @@ +// 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 item 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 +// flips 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 flips + // takenFlips indexes. + head *head + // takenFlips holds ascending positions in head.flips: the + // dependencies this path assumes the unpreferred way. Empty at the root. + takenFlips []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 most likely path. +func (n node) depth() int { return len(n.takenFlips) } + +// 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 flips 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 +// flips 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.takenFlips { + if a.takenFlips[n] != b.takenFlips[n] { + return a.takenFlips[n] < b.takenFlips[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 node. 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 most likely path in the queue. The RFC in +// doc/rfc/submitqueue/speculation-generator-best-first.md 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.takenFlips), + 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 +// flips; the root takes none, and a head with n flips 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 flipCost no greater +// than the one they replace or build on, since flips are sorted cheapest +// first and every flipCost is at most 0, and scoreFor sums both from +// bestScore 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 bestScore+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 flips are worked out when its root is expanded, which is + // this call if n is that root. + n.head.resolveFlips() + + taken := n.takenFlips + k := n.depth() + last := -1 + if k > 0 { + last = taken[k-1] + } + // Every flip 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.flips) { + return + } + + // Take one more flip, 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 flip 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 flips 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, takenFlips []int) { + heap.Push(&it.heap, node{ + head: h, + takenFlips: takenFlips, + score: h.scoreFor(takenFlips), + }) +} 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) +}