diff --git a/submitqueue/extension/scorer/README.md b/submitqueue/extension/scorer/README.md index 21c4ee60..0a7a33ad 100644 --- a/submitqueue/extension/scorer/README.md +++ b/submitqueue/extension/scorer/README.md @@ -1,64 +1,17 @@ -# Scorer +# scorer -Vendor-agnostic interface for computing success probability scores for code changes. +A `Scorer` returns the probability that a batch's build succeeds — a number between 0.0 and 1.0. It is handed the batch identity and resolves the batch's changes itself through an injected `changeset.Resolver`, so callers pass an `entity.Batch` and nothing more. -## Interface +Callers may score every batch a queue is waiting on, so implementations should be cheap. A speculation run scores each batch at most once, but it does not carry results across runs; anything expensive belongs behind the implementation's own cache. -### Scorer - -Computes a success probability for a given change. - -```go -type Scorer interface { - Score(ctx context.Context, change entity.Change) (float64, error) -} -``` - -- **change**: A `entity.Change` identifying the code change to score. -- **Score**: Returns a probability between 0.0 and 1.0 indicating the likelihood of a successful land. Returns an error if scoring fails. +Like the other extensions, a `Scorer` is selected **per queue** by the wiring layer through the `Config` (queue name) and `Factory` interface. ## Implementations -### Heuristic - -Scores a change by extracting a numeric value via a `ValueFunc` and matching it against ordered buckets. Each bucket maps a `[Min, Max]` range to a probability. - -```go -s := heuristic.New( - []heuristic.Bucket{ - {Min: 0, Max: 5, Score: 0.95}, - {Min: 6, Max: 20, Score: 0.75}, - {Min: 21, Max: 100, Score: 0.5}, - }, - func(ctx context.Context, change entity.Change) (int, error) { - // resolve the change into a numeric metric - return filesChanged, nil - }, -) - -score, err := s.Score(ctx, change) -``` - -### Composite - -Combines multiple named scorers into a single score using a reduce function. The reduce function receives a `map[string]float64` mapping scorer names to their scores, enabling domain-aware aggregation. - -Built-in reduce functions: `Min`, `Max`, `Avg`. - -```go -s := composite.New( - map[string]scorer.Scorer{ - "files": fileScorer, - "deps": depScorer, - }, - composite.Min, -) +**`heuristic`** scores a batch by extracting one number from its changes and matching that against ordered buckets, each mapping a `[Min, Max]` range to a probability. The extraction is a caller-supplied `ValueFunc` over the resolved `entity.BatchChanges`, so the same bucketing works for files touched, lines changed, or any other metric. -score, err := s.Score(ctx, change) -``` +**`composite`** runs several named scorers and reduces their scores to one. The reduce function receives the scores keyed by scorer name, so it can weigh sources differently rather than treating them as interchangeable; `Min`, `Max`, and `Avg` are provided. -## Implementing a Backend +## Adding a backend -1. Create `extension/scorer/{backend}/` directory -2. Implement the `Scorer` interface -3. Accept `entity.Change` and resolve it into whatever data the implementation needs +Create a package under `scorer//` whose `New(...)` returns a `scorer.Scorer`, injecting whatever it needs at construction — a `changeset.Resolver` to reach the batch's changes, a metrics scope, any client. Do not add a `Config` or `Factory` implementation here; per-queue routing and the factory adapter live in the wiring layer. diff --git a/submitqueue/extension/scorer/scorer.go b/submitqueue/extension/scorer/scorer.go index b3af1b09..da5ce729 100644 --- a/submitqueue/extension/scorer/scorer.go +++ b/submitqueue/extension/scorer/scorer.go @@ -22,11 +22,17 @@ import ( "github.com/uber/submitqueue/submitqueue/entity" ) -// Scorer computes a success probability score for a batch based on its changes. +// Scorer computes the probability that a batch's build succeeds, based on its +// changes. type Scorer interface { - // Score returns a probability between 0.0 and 1.0 indicating the likelihood - // of a successful land for the given batch. It is handed the batch identity - // and resolves the batch's changes itself through an injected changeset.Resolver. + // Score returns a probability between 0.0 and 1.0 that the given batch's + // build succeeds. It is handed the batch identity and resolves the batch's + // changes itself through an injected changeset.Resolver. + // + // Callers may score every batch a queue is waiting on, so implementations + // should be cheap: a speculation run scores each batch at most once, but it + // does not carry results over to the next run, so anything expensive to + // compute belongs behind the implementation's own cache. Score(ctx context.Context, batch entity.Batch) (float64, error) } diff --git a/submitqueue/extension/speculation/generator/BUILD.bazel b/submitqueue/extension/speculation/generator/BUILD.bazel new file mode 100644 index 00000000..2e290235 --- /dev/null +++ b/submitqueue/extension/speculation/generator/BUILD.bazel @@ -0,0 +1,9 @@ +load("@rules_go//go:def.bzl", "go_library") + +go_library( + name = "go_default_library", + srcs = ["generator.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/generator", + visibility = ["//visibility:public"], + deps = ["//submitqueue/entity:go_default_library"], +) diff --git a/submitqueue/extension/speculation/generator/README.md b/submitqueue/extension/speculation/generator/README.md new file mode 100644 index 00000000..f44f5852 --- /dev/null +++ b/submitqueue/extension/speculation/generator/README.md @@ -0,0 +1,9 @@ +# generator + +The `generator` package is a piece the `standard` `Speculator` is built from: a `Generator` produces the queue's candidate paths as one stream, best first, across all heads. It is **not** controller-facing — the speculate controller only knows the `Speculator` contract, and a different `Speculator` need not split its work this way. So there is no `Config` or `Factory` here; a `Generator` is chosen when the `standard` `Speculator` is constructed. + +`Open` starts the stream over the queue's live batches and returns a `PathIterator`. The caller pulls one candidate at a time and the generator does only the work that answer needs. A cancelled or expired context ends the stream with its error. + +Candidates never repeat and never contradict a known fact. Beyond that, the order is the `Generator`'s own: it yields candidates in whatever ranking it implements, and each carries the score it ranked by — higher first, on a scale the generator defines. Consumers take the iterator in the order given and do not interpret the score. Scores mean something only within the run and are never stored. + +The generator offers every path in the space, including paths whose builds already ran. Suppressing finished paths is the `Allocator`'s job, since that is the piece reconciling candidates against the stored path sets. diff --git a/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel new file mode 100644 index 00000000..ebdca570 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel @@ -0,0 +1,30 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = [ + "bestfirst.go", + "head.go", + "iterator.go", + ], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/generator/bestfirst", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/scorer:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = ["bestfirst_test.go"], + embed = [":go_default_library"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/scorer:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/generator/bestfirst/README.md b/submitqueue/extension/speculation/generator/bestfirst/README.md new file mode 100644 index 00000000..1e3a524b --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/README.md @@ -0,0 +1,47 @@ +# bestfirst + +A head waiting on unfinished dependencies has one path per combination of outcomes — 2ⁿ of them — while callers want the best few. `bestfirst` hands them out in score order without building the rest. + +Dependencies that already resolved carry a forced assumption: succeeded means the head builds on top of it, failed or cancelled means it builds without. They stay in the path but drop out of the search. Each unresolved one is a two-way assumption whose probability — that the dependency's build succeeds — comes from an injected `scorer.Scorer`, remembered per dependency so one shared by several heads is scored once. A dependency that is not a live batch cannot be scored at all, and gets a high default, since nearly every batch a queue accepts does build successfully. A path's score is the probability that all its assumptions hold. + +**The best path is not "assume everything succeeds."** Each dependency takes its *preferred* assumption — the one with the higher probability — so a dependency that will probably fail is assumed to fail. Every other path takes the unpreferred assumption on some of them, and each of those costs a known factor. + +**Scores are logarithms.** Scores use summed log probabilities to preserve ordering without underflow: a plain product across a wide head rounds to zero, ties, and loses the order entirely. A score is at most 0, and an assumption that cannot happen scores negative infinity. The package doc in `bestfirst.go` shows the failure it avoids. + +**Laziness.** One heap holds every path built but not yet handed out, across all heads. Each path handed out is replaced by the one or two that come after it within the same head, so the heap always holds each head's current best and the top of it is the best path in the queue. + +The full 2ⁿ path space is never generated, and no path is ever compared against more than the handful in the heap. Work is proportional to the number of heads plus the number of candidates pulled: `Open` scores each unresolved dependency once and sorts each head's alternatives, then every later path costs one build. What is *not* sorted is the path space itself — the alternatives within a head are, once, at `Open`. + +`bestfirst` discards nothing: pull long enough and every path for every head comes out. It does no conflict relaxation (`ignored` assumptions) — that is the `Speculator`'s call, not the generator's. + +## Worked example + +Say a head waits on three dependencies with scores d0 0.9, d1 0.8, d2 0.6. All three are better than even, so its best path assumes all three succeed, with probability 0.9 × 0.8 × 0.6 = 0.432. The tables below use probabilities because they read more easily; the code compares their logarithms. Sorted cheapest-first the alternatives are a0=d2, a1=d1, a2=d0. + +Every other path is the best path with some alternatives taken, written as the set of taken alternatives: `{a0}` takes only the cheapest, `{a0,a1}` the two cheapest, and so on. `expand` in `iterator.go` grows this space one step at a time — each path handed out queues the one or two that come just after it: + +``` +{} -> {0} +{0} -> {0,1}, {1} +{1} -> {1,2}, {2} +{0,1} -> {0,1,2}, {0,2} +{2}, {0,2}, {1,2}, {0,1,2} have nothing left to take above position 2 +``` + +Pulling all eight paths then goes: + +``` +handed out still waiting afterwards +{} .432 {a0}·.288 +{a0} .288 {a1}·.108 {a0,a1}·.072 +{a1} .108 {a0,a1}·.072 {a2}·.048 {a1,a2}·.012 +{a0,a1} .072 {a2}·.048 {a0,a2}·.032 {a1,a2}·.012 {a0,a1,a2}·.008 +{a2} .048 {a0,a2}·.032 {a1,a2}·.012 {a0,a1,a2}·.008 +{a0,a2} .032 {a1,a2}·.012 {a0,a1,a2}·.008 +{a1,a2} .012 {a0,a1,a2}·.008 +{a0,a1,a2} .008 — +``` + +The scores come out in order without the eight paths ever being enumerated or sorted together, and the waiting list never holds more than four of them. Stop after two pulls and only four paths were ever built: the head's best, plus the successors queued by those two. + +The proof that this walk reaches every path exactly once, in score order, is on `expand` in `iterator.go`. The behavior is pinned by tests in `bestfirst_test.go`. diff --git a/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go new file mode 100644 index 00000000..ce9580d1 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go @@ -0,0 +1,178 @@ +// 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. +// +// Three terms carry the package: +// +// - An assumption is what a path expects of one dependency — that it succeeds +// or that it does not. Once the dependency resolves the assumption is +// forced rather than chosen. Of the two, the one with the higher +// probability is preferred and the other unpreferred. +// - An alternative is the option to take the unpreferred assumption for a +// dependency still unresolved. The head's best path takes none of them; +// every other path takes some subset. +// - A candidate is one complete set of assumptions — a path — waiting in the +// heap with its score. +// +// # 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 comparing path IDs, 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. The full space is never enumerated: Open builds one path per +// head, and each pull builds at most the two that come after the path it hands +// out. Pulling k paths builds at most 1+2k of them, whatever n is. +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 ( + // scoreCutoff is where the preferred assumption flips: at or above it a + // dependency is assumed to succeed, below it assumed to fail. + // + // Not a tuning knob. The preferred side is whichever one this favors, so + // moving it would leave that side below probability 0.5 and relativeScore + // above 0, which breaks the best-first ordering. + scoreCutoff = 0.5 + + // defaultScore 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. + defaultScore = 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 works out each head's alternatives and queues that head's best path, +// ready to hand out. All the scoring happens here, which is why Next never +// calls the scorer. +// +// 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 + } + scores := newScoreCache(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, batchByID, scores) + if err != nil { + return nil, err + } + // Queue the head's best path — no alternatives taken, every dependency + // on its preferred assumption. It is the only path built up front; the + // rest grow out of it as the caller pulls. + it.push(h, []int{}) + } + return it, nil +} + +// scoreCache answers "what is this dependency's probability of building +// successfully", scoring each one at most once per run. Several heads usually +// wait on the same dependency, and there is no reason to score it once per +// head. +// +// The cache lives for one Open and no longer. A later run scores the same batch +// 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 scoreCache struct { + scorer scorer.Scorer + batches map[string]entity.Batch + scores map[string]float64 +} + +func newScoreCache(sc scorer.Scorer, batches map[string]entity.Batch) *scoreCache { + return &scoreCache{scorer: sc, batches: batches, scores: make(map[string]float64)} +} + +// of returns the given dependency's score: the probability its build succeeds. +// +// 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 *scoreCache) of(ctx context.Context, batchID string) (float64, error) { + if score, ok := c.scores[batchID]; ok { + return score, nil + } + batch, ok := c.batches[batchID] + if !ok { + // Not a live batch, so there is nothing to score. Remember the default + // too, so we don't look it up again. + c.scores[batchID] = defaultScore + return defaultScore, nil + } + score, err := c.scorer.Score(ctx, batch) + if err != nil { + return 0, err + } + if math.IsNaN(score) || score < 0 || score > 1 { + return 0, fmt.Errorf("scorer returned %v for batch %q: want a probability in [0, 1]", score, batchID) + } + c.scores[batchID] = score + return score, nil +} 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..314e30ad --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go @@ -0,0 +1,813 @@ +// 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 +} + +// iteratorOf reaches into the concrete iterator, so laziness invariants can be +// asserted on how many paths have actually been built. +func iteratorOf(t *testing.T, iter generator.PathIterator) *iterator { + t.Helper() + it, ok := iter.(*iterator) + require.True(t, ok, "iterator is not the bestfirst iterator") + return it +} + +// countingScorer records how many times each batch is scored. +type countingScorer struct { + scores map[string]float64 + calls map[string]int + total int +} + +func newCountingScorer(scores map[string]float64) *countingScorer { + return &countingScorer{scores: scores, calls: map[string]int{}} +} + +func (c *countingScorer) Score(_ context.Context, b entity.Batch) (float64, error) { + c.calls[b.ID]++ + c.total++ + if v, ok := c.scores[b.ID]; ok { + return v, nil + } + return 0.5, nil +} + +// errScorer always fails, to exercise error propagation from scoring. +type errScorer struct{} + +func (errScorer) Score(context.Context, entity.Batch) (float64, error) { + return 0, assert.AnError +} + +// constScorer scores every batch identically, regardless of ID. +type constScorer struct{ v float64 } + +func (c constScorer) Score(context.Context, entity.Batch) (float64, error) { return c.v, nil } + +// wideHead builds one Speculating head over n unresolved dependencies, each at a +// distinct score so no two combinations tie. +func wideHead(n int) ([]entity.Batch, scorer.Scorer) { + head := entity.Batch{ID: "q/head", State: entity.BatchStateSpeculating} + batches := []entity.Batch{} + scores := map[string]float64{} + for i := 0; i < n; i++ { + dep := fmt.Sprintf("q/d%02d", i) + head.Dependencies = append(head.Dependencies, dep) + // Created deps are unresolved (so they're scored) but not eligible heads. + batches = append(batches, entity.Batch{ID: dep, State: entity.BatchStateCreated}) + scores[dep] = 0.55 + 0.02*float64(i) + } + return append(batches, head), scored(scores) +} + +func TestBestFirst_OrderingAndEnumeration(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating}, + {ID: "q/3", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1", "q/2"}}, + } + sc := scored(map[string]float64{"q/1": 0.9, "q/2": 0.8}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/3") + + // Two unresolved dependencies -> 2^2 candidates. + require.Len(t, cands, 4) + + // Best-first: the all-succeeds path leads, scoring the product of the + // dependency scores (0.9 * 0.8), and scores descend from there. + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(cands[0].Path, "q/1")) + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(cands[0].Path, "q/2")) + assert.InDelta(t, math.Log(0.72), cands[0].RankingScore, 1e-9) + for i := 1; i < len(cands); i++ { + assert.LessOrEqual(t, cands[i].RankingScore, cands[i-1].RankingScore) + } + // The least optimistic candidate excludes both dependencies: (1-0.9)(1-0.8). + last := cands[len(cands)-1] + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(last.Path, "q/1")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(last.Path, "q/2")) + assert.InDelta(t, math.Log(0.02), last.RankingScore, 1e-9) +} + +func TestBestFirst_PinsResolvedDependencies(t *testing.T) { + // A resolved dependency is a fact: its assumption is forced by its + // state and it never varies across the head's paths. + tests := []struct { + name string + state entity.BatchState + want entity.DependencyAssumption + }{ + {name: "succeeded dependency forced to succeeds", state: entity.BatchStateSucceeded, want: entity.DependencyAssumptionSucceeds}, + {name: "failed dependency forced to fails", state: entity.BatchStateFailed, want: entity.DependencyAssumptionFails}, + {name: "cancelled dependency forced to fails", state: entity.BatchStateCancelled, want: entity.DependencyAssumptionFails}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: tt.state}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1"}}, + } + iter, err := New(scored(nil)).Open(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/2") + + // The pinned dependency has no alternative, so there is exactly one + // path, and it carries the forced assumption. + require.Len(t, cands, 1) + assert.Equal(t, tt.want, assumptionFor(cands[0].Path, "q/1")) + }) + } +} + +func TestBestFirst_ResolvedDependenciesDropOutOfSearch(t *testing.T) { + // Three dependencies, but two are resolved facts: only q/open is still + // unresolved, so the head yields 2 paths rather than 8. + batches := []entity.Batch{ + {ID: "q/succeeded", State: entity.BatchStateSucceeded}, + {ID: "q/failed", State: entity.BatchStateFailed}, + {ID: "q/open", State: entity.BatchStateCreated}, + {ID: "q/H", State: entity.BatchStateSpeculating, + Dependencies: []string{"q/succeeded", "q/failed", "q/open"}}, + } + iter, err := New(scored(map[string]float64{"q/open": 0.7})). + Open(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/H") + + require.Len(t, cands, 2) + // Pinned assumptions are identical on both paths and stay in queue order; the + // score reflects only the open dependency, since a fact is a certainty. + assert.Equal(t, "q/succeeded=succeeds;q/failed=fails;q/open=succeeds;", assumptionKey(cands[0].Path)) + assert.InDelta(t, math.Log(0.7), cands[0].RankingScore, 1e-9) + assert.Equal(t, "q/succeeded=succeeds;q/failed=fails;q/open=fails;", assumptionKey(cands[1].Path)) + assert.InDelta(t, math.Log(0.3), cands[1].RankingScore, 1e-9) +} + +func TestBestFirst_EmitsExactSequenceAcrossHeads(t *testing.T) { + // q/solo has nothing to wait on; q/pair assumes outcomes for two in-flight batches. The + // two heads interleave by score rather than draining one at a time. + batches := []entity.Batch{ + {ID: "q/d0", State: entity.BatchStateCreated}, + {ID: "q/d1", State: entity.BatchStateCreated}, + {ID: "q/solo", State: entity.BatchStateSpeculating}, + {ID: "q/pair", State: entity.BatchStateSpeculating, Dependencies: []string{"q/d0", "q/d1"}}, + } + sc := scored(map[string]float64{"q/d0": 0.9, "q/d1": 0.8}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + want := []struct { + head string + assumptions string + probability float64 + }{ + {head: "q/solo", assumptions: "", probability: 1.0}, + {head: "q/pair", assumptions: "q/d0=succeeds;q/d1=succeeds;", probability: 0.72}, + {head: "q/pair", assumptions: "q/d0=succeeds;q/d1=fails;", probability: 0.18}, + {head: "q/pair", assumptions: "q/d0=fails;q/d1=succeeds;", probability: 0.08}, + {head: "q/pair", assumptions: "q/d0=fails;q/d1=fails;", probability: 0.02}, + } + require.Len(t, cands, len(want)) + for i, w := range want { + assert.Equal(t, w.head, cands[i].Path.Head, "head at %d", i) + assert.Equal(t, w.assumptions, assumptionKey(cands[i].Path), "assumptions at %d", i) + assert.InDelta(t, math.Log(w.probability), cands[i].RankingScore, 1e-9, "score at %d", i) + } +} + +func TestBestFirst_PreferredAssumptionFollowsScore(t *testing.T) { + // q/low is more likely to fail than to succeed, so the head's best path assumes + // it fails. The leading path is the preferred-assumption path, not the + // all-succeeds one. + batches := []entity.Batch{ + {ID: "q/high", State: entity.BatchStateCreated}, + {ID: "q/low", State: entity.BatchStateCreated}, + {ID: "q/H", State: entity.BatchStateSpeculating, + Dependencies: []string{"q/high", "q/low"}}, + } + sc := scored(map[string]float64{"q/high": 0.8, "q/low": 0.3}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + want := []struct { + assumptions string + probability float64 + }{ + {assumptions: "q/high=succeeds;q/low=fails;", probability: 0.56}, + {assumptions: "q/high=succeeds;q/low=succeeds;", probability: 0.24}, + {assumptions: "q/high=fails;q/low=fails;", probability: 0.14}, + {assumptions: "q/high=fails;q/low=succeeds;", probability: 0.06}, + } + require.Len(t, cands, len(want)) + for i, w := range want { + assert.Equal(t, w.assumptions, assumptionKey(cands[i].Path), "assumptions at %d", i) + assert.InDelta(t, math.Log(w.probability), cands[i].RankingScore, 1e-9, "score at %d", i) + } + + // Spelled out: the all-succeeds path exists but only ranks second. + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(cands[0].Path, "q/low")) +} + +func TestBestFirst_OnlySpeculatingHeadsProduceCandidates(t *testing.T) { + // Batches in any state supply facts about their dependents, but only a + // Speculating batch is a head worth proposing work on. + tests := []struct { + state entity.BatchState + want bool + }{ + {state: entity.BatchStateUnknown}, + {state: entity.BatchStateCreated}, + {state: entity.BatchStateSpeculating, want: true}, + {state: entity.BatchStateMerging}, + {state: entity.BatchStateSucceeded}, + {state: entity.BatchStateFailed}, + {state: entity.BatchStateCancelling}, + {state: entity.BatchStateCancelled}, + } + + for _, tt := range tests { + name := string(tt.state) + if name == "" { + name = "unknown" + } + t.Run(name, func(t *testing.T) { + batches := []entity.Batch{{ID: "q/1", State: tt.state}} + + iter, err := New(scored(nil)).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + if !tt.want { + assert.Empty(t, cands) + return + } + require.Len(t, cands, 1) + assert.Equal(t, "q/1", cands[0].Path.Head) + }) + } +} + +func TestBestFirst_HeadWithNoDependencies(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + } + + iter, err := New(scored(nil)).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + require.Len(t, cands, 1) + assert.Equal(t, "q/1", cands[0].Path.Head) + assert.Empty(t, cands[0].Path.Dependencies) + assert.InDelta(t, math.Log(1.0), cands[0].RankingScore, 1e-9) +} + +func TestBestFirst_PropagatesScorerError(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/1", State: entity.BatchStateSpeculating}, + {ID: "q/2", State: entity.BatchStateSpeculating, Dependencies: []string{"q/1"}}, + } + + iter, err := New(errScorer{}).Open(context.Background(), batches) + assert.Error(t, err) + assert.Nil(t, iter) +} + +func TestBestFirst_AbsentDependencyUsesDefaultScore(t *testing.T) { + batches := []entity.Batch{ + {ID: "q/H", State: entity.BatchStateSpeculating, Dependencies: []string{"q/missing"}}, + } + + // q/missing is not among the live batches, so it cannot be scored at all — + // the generator uses the default score without consulting the + // scorer, whatever the scorer would have said. + iter, err := New(constScorer{0.1}).Open(context.Background(), batches) + require.NoError(t, err) + cands := forHead(drainAll(t, iter), "q/H") + + // Assumed almost certain to succeed, so the best path assumes it does. + require.Len(t, cands, 2) + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(cands[0].Path, "q/missing")) + assert.InDelta(t, math.Log(defaultScore), cands[0].RankingScore, 1e-9) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(cands[1].Path, "q/missing")) + assert.InDelta(t, math.Log(1-defaultScore), cands[1].RankingScore, 1e-9) +} + +func TestBestFirst_GeneratesOnlyWhatIsPulled(t *testing.T) { + // 12 unresolved dependencies is an outcome space of 4096 paths. + const deps, space = 12, 1 << 12 + batches, sc := wideHead(deps) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + it := iteratorOf(t, iter) + + // Everything ever built is either handed out already or still waiting in + // the heap. Open builds the head's best path and nothing more. + require.Equal(t, 1, it.heap.Len()) + + pulled := 0 + for i := 0; i < 3; i++ { + _, ok, err := iter.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + pulled++ + } + + // Each pull builds at most two follow-ups, so three pulls can have built at + // most 1+2*3 paths — a tiny fraction of the space. + built := pulled + it.heap.Len() + assert.LessOrEqual(t, built, 7) + assert.Less(t, built, space) +} + +func TestBestFirst_DrainYieldsEveryCombinationOnce(t *testing.T) { + deps := []string{"q/d0", "q/d1", "q/d2"} + batches := []entity.Batch{ + {ID: "q/d0", State: entity.BatchStateCreated}, + {ID: "q/d1", State: entity.BatchStateCreated}, + {ID: "q/d2", State: entity.BatchStateCreated}, + {ID: "q/head", State: entity.BatchStateSpeculating, Dependencies: deps}, + } + sc := scored(map[string]float64{"q/d0": 0.9, "q/d1": 0.7, "q/d2": 0.6}) + + iter, err := New(sc).Open(context.Background(), batches) + require.NoError(t, err) + cands := drainAll(t, iter) + + require.Len(t, cands, 8) + + // Every include/exclude combination appears exactly once, and each path + // carries its assumptions in queue order. + seen := map[string]int{} + for _, c := range cands { + got := make([]string, 0, len(c.Path.Dependencies)) + for _, dep := range c.Path.Dependencies { + got = append(got, dep.Batch) + } + assert.Equal(t, deps, got, "assumptions must stay in dependency order") + seen[assumptionKey(c.Path)]++ + } + for mask := 0; mask < 8; mask++ { + var want strings.Builder + for i, dep := range deps { + assumption := entity.DependencyAssumptionFails + if mask&(1< 0 { + assert.LessOrEqual(t, cands[i].RankingScore, cands[i-1].RankingScore) + } + } +} diff --git a/submitqueue/extension/speculation/generator/bestfirst/head.go b/submitqueue/extension/speculation/generator/bestfirst/head.go new file mode 100644 index 00000000..54726e30 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/head.go @@ -0,0 +1,206 @@ +// 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" + "math" + "sort" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// alternative is the option to take the unpreferred assumption for one +// unresolved dependency. +// +// Each unresolved dependency has a preferred assumption — the one with the +// higher probability, which the head's best path takes — and this alternative, +// the unpreferred one. Every other path takes some subset of the head's +// alternatives, each adding its relativeScore to the path's score. +// +// An alternative holds no taken/untaken state: each path takes its own subset, +// readable off the path's assumptions. +type alternative struct { + // at is the index into the head's bestPath.Dependencies at which this + // alternative swaps the assumption. The alternatives slice is sorted by + // relativeScore, which scrambles queue order, so each alternative carries + // the index it came from. + at int + // assumption is what a path taking this alternative assumes — the + // unpreferred side of the dependency at index at. + assumption entity.DependencyAssumption + // relativeScore is what taking this alternative adds to a path's score: + // log(unpreferred probability) - log(preferred probability). Always at most + // 0, since the unpreferred side never has the higher probability, and + // negative infinity when the unpreferred side cannot happen. + // + // A difference, not just the unpreferred side's log, because bestPathScore + // already counts the preferred one: with preferred 0.8 and unpreferred 0.2, + // swapping 0.8 out of the sum for 0.2 is one add of log(0.2)-log(0.8). + relativeScore float64 +} + +// head is one batch we might build, and everything needed to enumerate its +// paths. It holds no heap of its own — every head's paths compete in the +// iterator's single heap. +type head struct { + // bestPath is the path this head would build if every unresolved dependency + // took its preferred assumption: the head batch plus one assumption per + // dependency, in queue order, forced where the dependency has already + // resolved. Every other path is this one with some alternatives applied. + bestPath entity.SpeculationPath + // alternatives is one entry per unresolved dependency — the option to swap + // that dependency to its unpreferred assumption — sorted cheapest first. + alternatives []alternative + // bestPathScore is bestPath's score: the preferred probabilities' + // logarithms summed. Every other path's score is bestPathScore plus the + // relativeScore of each alternative it takes. Always finite — every + // preferred probability is at least 0.5, so none of the logs is infinite. + bestPathScore float64 +} + +// newHead works out a head's search space. Dependencies that already resolved +// contribute a forced assumption and drop out of the search; each of the rest +// becomes an alternative, ordered so the cheapest one comes first. +// +// A cancelled or expired ctx aborts with its error: a head can carry hundreds +// of dependencies, so the check is per dependency rather than per head. +func newHead( + ctx context.Context, + batch entity.Batch, + batchByID map[string]entity.Batch, + scores *scoreCache, +) (*head, error) { + dependencies := make([]entity.PathDependency, len(batch.Dependencies)) + alternatives := make([]alternative, 0, len(batch.Dependencies)) + bestPathScore := 0.0 + for i, dep := range batch.Dependencies { + if err := ctx.Err(); err != nil { + return nil, err + } + if assumption, settled := settledAssumption(batchByID[dep].State); settled { + dependencies[i] = entity.PathDependency{Batch: dep, Assumption: assumption} + continue + } + score, err := scores.of(ctx, dep) + if err != nil { + return nil, err + } + // Prefer the assumption with the higher probability; on an exact tie, + // prefer succeeds. preferredScore is that side's probability, so it is + // always at least 0.5 and 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. + preferredScore, unpreferredScore := score, 1-score + preferredAssumption, unpreferredAssumption := entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails + if score < scoreCutoff { + preferredScore, unpreferredScore = 1-score, score + preferredAssumption, unpreferredAssumption = entity.DependencyAssumptionFails, entity.DependencyAssumptionSucceeds + } + + dependencies[i] = entity.PathDependency{Batch: dep, Assumption: preferredAssumption} + bestPathScore += math.Log(preferredScore) + alternatives = append(alternatives, alternative{ + at: i, + assumption: unpreferredAssumption, + // math.Log(0) is -Inf, which is exactly right for an alternative + // that cannot happen: it sinks every path taking it below every + // path that does not, and stays -Inf however much is added. + relativeScore: math.Log(unpreferredScore) - math.Log(preferredScore), + }) + } + // Cheapest alternative first — the assumption we'd mind changing least, so + // the one whose relativeScore is closest to 0. expand depends on this + // order: it only ever moves a taken alternative to the next position along, + // so having them sorted is what guarantees a path never scores above the one + // it came from. + sort.Slice(alternatives, func(i, j int) bool { + if alternatives[i].relativeScore != alternatives[j].relativeScore { + return alternatives[i].relativeScore > alternatives[j].relativeScore + } + return alternatives[i].at < alternatives[j].at + }) + + return &head{ + bestPath: entity.SpeculationPath{Head: batch.ID, Dependencies: dependencies}, + alternatives: alternatives, + bestPathScore: bestPathScore, + }, nil +} + +// settledAssumption 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 settle. Cancelling deliberately does not: cancellation +// is best-effort, so a cancelling batch can still succeed. Settling it on 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 settledAssumption(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 + } +} + +// pathFor builds the path that takes the given alternatives, and its score. +// takenAlternatives holds ascending positions in h.alternatives. It starts from +// h.bestPath, swaps each taken alternative's dependency to that alternative's +// assumption, and adds its relativeScore. Dependencies that already resolved +// are certainties, so they have no alternative and never change the score. +// +// It is the inverse of takenAlternatives: h.pathFor(h.takenAlternatives(p)) +// reconstructs p. +func (h *head) pathFor(takenAlternatives []int) (entity.SpeculationPath, float64) { + // Work on a copy. h.bestPath is the template every one of this head's paths + // starts from, so writing an assumption straight into it would corrupt every + // path built after this one. + dependencies := make([]entity.PathDependency, len(h.bestPath.Dependencies)) + copy(dependencies, h.bestPath.Dependencies) + + pathScore := h.bestPathScore + for _, i := range takenAlternatives { + a := h.alternatives[i] + dependencies[a.at].Assumption = a.assumption + pathScore += a.relativeScore + } + return entity.SpeculationPath{Head: h.bestPath.Head, Dependencies: dependencies}, pathScore +} + +// takenAlternatives reads back which of h's alternatives a path takes, as +// ascending positions in h.alternatives: exactly the ones whose dependency the +// path assumes the unpreferred way. It inverts pathFor, so the iterator needs +// no bookkeeping on the paths it holds. +func (h *head) takenAlternatives(path entity.SpeculationPath) []int { + var taken []int + for i, a := range h.alternatives { + if path.Dependencies[a.at].Assumption == a.assumption { + taken = append(taken, i) + } + } + return taken +} diff --git a/submitqueue/extension/speculation/generator/bestfirst/iterator.go b/submitqueue/extension/speculation/generator/bestfirst/iterator.go new file mode 100644 index 00000000..2793c20b --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/iterator.go @@ -0,0 +1,156 @@ +// 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" +) + +// candidate is one path waiting in the heap, built but not yet handed out. It +// remembers which head it belongs to so expand can build the paths that come +// after it; everything else about the path is readable off the path itself. +type candidate struct { + entity.CandidatePath + // pathID is Path.ID(), kept here so comparing two entries needn't rehash. + pathID string + head *head +} + +// candidateHeap keeps candidates best-first: highest score first, and when two +// scores tie, lowest path ID 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. +// +// The tie-break makes a run repeatable, not globally ID-sorted: only candidates +// in the heap at the same moment are compared. Two equal-scored paths that are +// never resident together come out in whatever order the walk builds them — +// always the same one, run to run. +type candidateHeap []candidate + +func (h candidateHeap) Len() int { return len(h) } +func (h candidateHeap) Less(i, j int) bool { + // Greater-than, not less-than: that is what makes this a max-heap, so the + // highest-scoring path ends up on top. + if h[i].RankingScore != h[j].RankingScore { + return h[i].RankingScore > h[j].RankingScore + } + return h[i].pathID < h[j].pathID +} +func (h candidateHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } +func (h *candidateHeap) Push(x any) { *h = append(*h, x.(candidate)) } +func (h *candidateHeap) Pop() any { + // Careful: this does not pick the best entry. container/heap has already + // swapped the best one to the end of the slice before calling us, so all we + // do is chop the last element off and hand it back. + old := *h + n := len(old) + e := old[n-1] + *h = old[:n-1] + return e +} + +// iterator hands out the queue's paths, best first, building them as it goes. +// +// One heap holds every path that has been built but not yet handed out, across +// all heads. Each time a path is handed out, the one or two paths that come +// just after it — within that same head — get built and go into the heap. So +// the heap always holds each head's current best, and taking the top of it is +// taking the best path in the queue. The README works through a full example. +type iterator struct { + heap candidateHeap +} + +// Next hands out the best remaining path across all heads. On the way out it +// builds whatever comes after that path, so the heap always has the next +// candidates ready. It returns ok=false once nothing is left to offer. +// +// 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 + } + c := heap.Pop(&it.heap).(candidate) + it.expand(c) + return c.CandidatePath, true, nil +} + +// expand builds the paths that come just after this one, within the same head. +// +// Describe a path by which alternatives it takes — ascending positions in the +// head's cheapest-first alternatives, read back off the path's assumptions. A +// head with n alternatives has 2^n such sets, 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 to move 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 set is built 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 +// set has exactly one parent, and from the best path every set is reachable by +// re-tracing those moves. That is why nothing needs to remember what has +// already been built. +// +// A path never outscores the one it came from. Both moves add a relativeScore +// no greater than the one they replace or build on (alternatives are sorted +// cheapest first, and every relativeScore is at most 0), and they add it to the +// identical computed prefix — so IEEE addition preserves the ordering as +// computed, not just in exact arithmetic. That is what lets the heap hand paths +// out in score order. +func (it *iterator) expand(c candidate) { + taken := c.head.takenAlternatives(c.Path) + k := len(taken) + last := -1 + if k > 0 { + last = taken[k-1] + } + // Every alternative past the last taken one is already used up, so there is + // no move left to make and this path has no follow-ups. + if last+1 >= len(c.head.alternatives) { + return + } + + // Take one more alternative, keeping the ones already taken. + extended := make([]int, k+1) + copy(extended, taken) + extended[k] = last + 1 + it.push(c.head, extended) + + // Move the last taken alternative one position along. The best path has + // none to move, so this second move doesn't apply to it. + if k > 0 { + advanced := make([]int, k) + copy(advanced, taken) + advanced[k-1] = last + 1 + it.push(c.head, advanced) + } +} + +// push builds the path that takes the given alternatives and puts it in the heap. +func (it *iterator) push(h *head, takenAlternatives []int) { + path, score := h.pathFor(takenAlternatives) + heap.Push(&it.heap, candidate{ + CandidatePath: entity.CandidatePath{Path: path, RankingScore: score}, + // Hash the path once here rather than every time the heap compares it. + pathID: path.ID(), + head: h, + }) +} diff --git a/submitqueue/extension/speculation/generator/generator.go b/submitqueue/extension/speculation/generator/generator.go new file mode 100644 index 00000000..42cebc17 --- /dev/null +++ b/submitqueue/extension/speculation/generator/generator.go @@ -0,0 +1,51 @@ +// 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. + Open(ctx context.Context, batches []entity.Batch) (PathIterator, error) +} + +// PathIterator is a pull-based stream of candidate paths. 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) +}