From 599849aed8dd0eea4a38074c1e6fa659997490b9 Mon Sep 17 00:00:00 2001 From: sergeyb Date: Wed, 5 Aug 2026 21:40:37 +0000 Subject: [PATCH] feat(speculation): add best-first speculation path generator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the Generator seam of the standard Speculator (see doc/rfc/submitqueue/speculation.md): a pull-based stream of candidate speculation paths over one queue snapshot, consumed by the Allocator. - generator: the Generator/PathIterator contract — lazy, snapshot-scoped, ctx-aborted; candidates never repeat and never contradict resolved facts. - generator/bestfirst: prices every path by the probability that all of its dependency assumptions hold — resolved deps pinned as facts, Merging/Cancelling priced as modal certainties, undecided deps scored via the injected scorer (once per batch per run) — and yields candidates in exact non-increasing price order through lazy add/shift enumeration merged across heads by one max-heap. Dependencies are normalized into canonical queue order so path IDs stay stable across runs; ties are deterministic. - doc/rfc/submitqueue/speculation-generator.md: design doc with the pricing model, enumeration scheme, a step-by-step two-run worked example, and future refinements (relaxation, unblocking weight, sensitivity pruning). Also regenerates storage/mock/request_batch_store_mock.go, whose committed header predated generation via 'make mocks' and kept check-mocks red. Co-Authored-By: Claude Fable 5 --- doc/rfc/index.md | 1 + doc/rfc/submitqueue/speculation-generator.md | 121 ++++++ .../speculation/generator/BUILD.bazel | 9 + .../extension/speculation/generator/README.md | 11 + .../generator/bestfirst/BUILD.bazel | 29 ++ .../speculation/generator/bestfirst/README.md | 21 + .../generator/bestfirst/bestfirst.go | 230 ++++++++++ .../generator/bestfirst/bestfirst_test.go | 396 ++++++++++++++++++ .../generator/bestfirst/iterator.go | 198 +++++++++ .../speculation/generator/generator.go | 59 +++ .../speculation/generator/mock/BUILD.bazel | 13 + .../generator/mock/generator_mock.go | 98 +++++ .../storage/mock/request_batch_store_mock.go | 4 +- 13 files changed, 1188 insertions(+), 2 deletions(-) create mode 100644 doc/rfc/submitqueue/speculation-generator.md create mode 100644 submitqueue/extension/speculation/generator/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/README.md create mode 100644 submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/bestfirst/README.md create mode 100644 submitqueue/extension/speculation/generator/bestfirst/bestfirst.go create mode 100644 submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go create mode 100644 submitqueue/extension/speculation/generator/bestfirst/iterator.go create mode 100644 submitqueue/extension/speculation/generator/generator.go create mode 100644 submitqueue/extension/speculation/generator/mock/BUILD.bazel create mode 100644 submitqueue/extension/speculation/generator/mock/generator_mock.go diff --git a/doc/rfc/index.md b/doc/rfc/index.md index 37e0d867..b4e0183a 100644 --- a/doc/rfc/index.md +++ b/doc/rfc/index.md @@ -18,6 +18,7 @@ Design documents and technical proposals, grouped by scope. Shared/cross-cutting - [Extension Contract](submitqueue/extension-contract.md) - When extensions take orchestrator identity (request/batch) and resolve granular content themselves vs. take controller-resolved data; revises the BuildRunner base/head contract - [Gateway Status and List APIs](submitqueue/status-list-api.md) - Gateway-owned request context, materialized current status, sqid or change-URI status lookup, and queue admission listing - [Speculation](submitqueue/speculation.md) - Why SubmitQueue speculates, the path/tree model, and the two pluggable seams: speculation-tree enumeration and path selection +- [Best-First Speculation Generator](submitqueue/speculation-generator.md) - Default Generator for the standard Speculator: prices every path by the probability its assumptions hold and streams candidates lazily in exactly that order (facts pinned, scorer-backed bets, add/shift best-first enumeration) - [Modular Queue Wiring](submitqueue/modular-queue-wiring.md) - Declare-don't-assemble engine (`pipeline.Construct`) that unifies topic registry, controller registration, DLQ pairing, and lifecycle ordering into one typed call; services self-declare via Deps struct + Stages slice, hosts own per-queue profiles and transport ## Stovepipe diff --git a/doc/rfc/submitqueue/speculation-generator.md b/doc/rfc/submitqueue/speculation-generator.md new file mode 100644 index 00000000..f245619a --- /dev/null +++ b/doc/rfc/submitqueue/speculation-generator.md @@ -0,0 +1,121 @@ +# Best-First Speculation Generator + +Design of `bestfirst`, the default Generator inside the standard Speculator (see [Speculation](speculation.md)). The Generator's job: given one snapshot of a queue's batches, stream candidate speculation paths in the order most likely to pay off, lazily, so the Allocator can spend the build budget off the top of the stream. + +## Model: every path is a bet + +A candidate path is a head in `BatchStateSpeculating` plus one assumption — *succeeds* or *fails* — per dependency. Building a path pays off only if every assumption matches how that dependency actually resolves; a build on a broken assumption is refuted by the controller and its CI time is wasted. So the natural ranking of candidates is the probability that the whole bet holds: + +``` +price(path) = ∏ over dependencies d: P(d succeeds) if the path assumes d succeeds + 1 − P(d succeeds) if the path assumes d fails +``` + +Dependency probabilities come from the snapshot state, with the injected `scorer.Scorer` filling in the genuinely undecided ones: + +| dependency state | P(succeeds) | kind | effect | +|---|---|---|---| +| Succeeded | 1 | fact | pinned *succeeds*; the *fails* branch is never generated | +| Failed, Cancelled | 0 | fact | pinned *fails*; the *succeeds* branch is never generated | +| Merging | 1 | modal bet | build already passed and hand-off is in flight; opposite branch suppressed | +| Cancelling | 0 | modal bet | halted, only a lost cancel race resurrects it; opposite branch suppressed | +| Created, Speculating | scorer | bet | both branches generated, priced p and 1−p | +| anything else / missing | — | incoherent | the head yields nothing this run (see below) | + +Facts honor the Generator contract that no candidate contradicts a resolved outcome. The two modal bets are a pricing policy, not facts: if the long shot lands anyway — a Merging batch whose merge fails, a Cancelling batch that merges because the cancel lost the race — the controller refutes the affected paths and the next run re-plans against the new facts, exactly as with any lost bet. Scorer probabilities of exactly 0 or 1 pin the same way facts do. Scores are fetched at `Open`, at most once per batch per run. + +Two properties fall out of pure price ordering, with no special cases: + +- **Sure work first.** A head whose dependencies are all decided has exactly one coherent path, priced 1, so real builds always outrank speculation. In a queue with no useful scores (everything near 0.5), deep speculation prices as 0.5^k and sinks — the stream degrades gracefully toward plain dependency-order building. +- **Hedging is automatic.** A dependency near 0.5 puts both of its branches adjacent in the stream. The branches are mutually exclusive, so funding both covers the head against either outcome — the classic small-change-behind-a-big-one bypass builds on the same mechanism (cover the whole outcome space, merge early), and here it simply falls out of the budget reaching that deep. + +## Lazy best-first enumeration + +A head with k undecided dependencies has 2^k coherent paths; they are never materialized. Per head: + +1. The **modal path** takes every undecided dependency at its more likely outcome. Its price is `P₀ = ∏ q_d` where `q_d = max(p_d, 1−p_d)` (ties at 0.5 go to *succeeds*); pinned dependencies contribute certainty and no factor. +2. Every other path is the modal path with a subset S of dependencies flipped: `price(S) = P₀ · ∏_{d∈S} r_d` with flip ratio `r_d = (1−q_d)/q_d ∈ (0,1]`. +3. Flips are sorted by descending ratio (cheapest flip first; queue order on ties) and subsets are enumerated through the **add/shift tree**: the children of a subset whose largest flip index is m are S∪{m+1} (*add*) and S\{m}∪{m+1} (*shift*). Every subset is generated exactly once, and each child's price is at most its parent's. + +A single max-heap merges all heads: seed it with every head's modal path, then pop the best entry, materialize it, and push its at most two children. Since children never outrank parents, popped prices are globally non-increasing — the stream is exactly best-first, at O(log n) per pull, with memory proportional to what was pulled. Ties break deterministically: fewer flips first (closer to the modal path), then insertion order; heads are seeded in queue order, so equally priced sure paths stream oldest first. + +Example enumeration for one head with undecided dependencies D (q=0.7, r≈0.43) and C (q=0.8, r=0.25), pinned prefix omitted, P₀ = 0.56: + +``` +∅ ──────────── price 0.56 {C:succeeds, D:succeeds} (modal) +└─ {D} ──────── price 0.24 {C:succeeds, D:fails} (add: ×0.43) + ├─ {D,C} ─── price 0.06 {C:fails, D:fails} (add: ×0.25) + └─ {C} ───── price 0.14 {C:fails, D:succeeds} (shift: ×0.25/0.43) +``` + +The tree only guarantees children never beat parents; the heap decides actual emission order (here ∅, {D}, {C}, {D,C}). + +## Run scope: the snapshot replaces the bookkeeping + +The abstract version of this algorithm needs machinery for a long-lived iterator: result caching, refutation cascades, entry invalidation when a dependency resolves mid-stream. None of that lives here, because the speculate controller reruns from scratch on every dirty signal: each run reads a fresh snapshot, opens a fresh stream, and persists only what it funds. Resolution "cascades" happen for free — the next snapshot simply pins the resolved dependency, every surviving path's price is implicitly reconditioned, and contradicted paths stop being generated at all. Duplicate-work suppression is the Allocator's half: it matches candidates to existing path sets by path ID, keeps in-flight paths in their slots, and skips terminal ones. + +That division makes path identity the load-bearing contract. `entity.SpeculationPath` hashes the head and the *ordered* dependency assumptions, so the Generator must emit dependencies in a canonical order or the same logical path would get a new ID every run and be rebuilt. The Generator normalizes dependencies into queue order — ascending batch counter from the documented `/batch/` ID format, with a total deterministic fallback for unparsable IDs — and collapses duplicate entries. Given equal snapshots, the whole stream is deterministic regardless of input slice order. + +A head whose snapshot cannot support coherent candidates — a dependency missing from the snapshot, listed as its own dependency, or in a state outside the table above — yields nothing this run. Skipping is safe because runs are cheap and self-healing: the head is re-planned the moment a complete snapshot arrives, and proposing nothing is always a valid Generator output. + +## Worked example + +Queue `q`, five batches, all Speculating. Edges point dependency → dependent; scorer probabilities in parentheses. + +``` + b1 (0.9) b2 (0.5) ← no dependencies + \ / \ + v v v + b3 (0.8) b4 (0.7) + \ / + v v + b5 (deps: b1, b2, b3, b4) +``` + +### Run 1: everything undecided + +Seeds, per head: b1 and b2 price 1 (no dependencies); b4 prices 0.5 (one coin-flip dependency); b3 prices 0.9×0.5 = 0.45; b5 prices 0.9×0.5×0.8×0.7 = 0.252. The stream begins: + +| # | head | assumptions (S=succeeds, F=fails) | price | note | +|---|------|-----------------------------------|-------|------| +| 1 | b1 | — | 1.00 | sure build | +| 2 | b2 | — | 1.00 | sure build | +| 3 | b4 | b2:S | 0.50 | best bet | +| 4 | b4 | b2:F | 0.50 | its hedge — b4 now covered either way | +| 5 | b3 | b1:S b2:S | 0.45 | | +| 6 | b3 | b1:S b2:F | 0.45 | coin-flip hedge again | +| 7 | b5 | b1:S b2:S b3:S b4:S | 0.252 | | +| 8 | b5 | b1:S b2:F b3:S b4:S | 0.252 | | +| 9 | b5 | b1:S b2:S b3:S b4:F | 0.108 | flips get deeper | +| 10 | b5 | b1:S b2:F b3:S b4:F | 0.108 | | + +An Allocator with budget 6 funds rows 1–6: both sure builds, and both branches of b4 and b3 across the b2 coin flip. The stream continues lazily (24 coherent paths total) only if pulled. + +### Run 2: b1 succeeded, b2 failed + +The next dirty signals bring a snapshot where b1 is Succeeded and b2 is Failed. No entries are repriced and nothing is cancelled *by the Generator* — the controller has already refuted the paths that assumed b2 succeeds, and the fresh stream simply prices the new facts in: + +| # | head | assumptions | price | note | +|---|------|-------------|-------|------| +| 1 | b3 | b1:S b2:F | 1.00 | was row 6 of run 1; same path ID, so the Allocator keeps its slot | +| 2 | b4 | b2:F | 1.00 | was the row-4 hedge; now the sure build — already in flight if funded in run 1 | +| 3 | b5 | b1:S b2:F b3:S b4:S | 0.56 | conditioned from 0.252: the b1 and b2 factors became certainty | +| 4 | b5 | b1:S b2:F b3:S b4:F | 0.24 | | +| 5 | b5 | b1:S b2:F b3:F b4:S | 0.14 | | +| 6 | b5 | b1:S b2:F b3:F b4:F | 0.06 | | + +b5's space collapsed from 16 paths to 4 — paths contradicting the b1/b2 facts are never generated — and every surviving path kept its ID, so nothing already built is rebuilt. This is the whole reconciliation story: pin, reprice, re-rank, all implicit in re-opening the stream. + +## Design notes and rejected alternatives + +- **Score the path's own head too?** No: the head's probability of passing does not change whether the *bet on its dependencies* pays off, and a passed build is informative even for a head that will fail. Head-quality weighting belongs to ranking policy evolution (below), not the payoff model. +- **A ranking-score floor (don't propose paths below price X)?** Rejected by the Speculation RFC: budget is the only rationing lever. A cheap hedge is worth a slot that would otherwise idle; the Allocator decides that, not the Generator. +- **Cache scorer results across runs?** The scorer contract already places caching behind the scorer's own interface; the Generator memoizes per run only, keeping runs stateless and reproducible. +- **Metrics/logging in the Generator?** Omitted: `Open` is pure CPU over one snapshot plus scorer calls, and both neighbors (scorer implementations, the speculate controller) already instrument their halves. + +## Future refinements + +- **Relaxation.** The path model supports *ignored* assumptions (see Conflict relaxation in the Speculation RFC); this Generator does not emit them yet. A relaxation policy slots in as a pre-pass that drops weak dependencies from the flip universe and marks them ignored on the template. +- **Unblocking weight.** Pure price ordering is myopic about information value: resolving a batch that many heads depend on collapses more of the space. A weight like 1 + λ·(unresolved dependents) multiplied into a head's prices would bias the stream toward unblocking without changing the machinery. +- **Sharper priors.** Passed or failed builds of a head's *other* paths are evidence about its dependencies-independent quality; a scorer (or a wrapper) reading recent path outcomes could sharpen probabilities between runs without touching the Generator. +- **Sensitivity pruning.** If conflict analysis can certify that a head's build outcome is independent of one dependency's content, that dependency needs no branch at all — each certified dependency halves a head's path space. This is the strongest practical lever for large closures and composes with relaxation. 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..53388421 --- /dev/null +++ b/submitqueue/extension/speculation/generator/README.md @@ -0,0 +1,11 @@ +# Speculation Generator + +The candidate-stream composition point inside the standard Speculator. A Generator turns one snapshot of a queue's batches into a pull-based stream of candidate speculation paths — each an `entity.CandidatePath`: a head in `BatchStateSpeculating`, one assumption per dependency in queue order, and the transient ranking score the stream is ordered by. + +This is not a controller-facing extension. The speculate controller depends only on the Speculator contract; the default Speculator opens a Generator's stream and hands it to an Allocator, and an alternate Speculator need not use or expose this seam. See the [Speculation RFC](../../../../doc/rfc/submitqueue/speculation.md) for the composition and the [Best-First Speculation Generator RFC](../../../../doc/rfc/submitqueue/speculation-generator.md) for the default ranking policy. + +The stream is lazy: beyond what ranking requires up front, a Generator computes only what the consumer pulls, so the path space — up to two branches per undecided dependency per head — is never materialized. Candidates never repeat and never contradict a resolved fact: a dependency that already succeeded is never assumed to fail, and one that already failed or was cancelled is never assumed to succeed. The batches slice is a caller-owned snapshot the Generator may retain but never mutates; a queue that has moved on is a new snapshot and a new stream. Both `Open` and `Next` abort on a cancelled context. + +## Implementations + +- [bestfirst](bestfirst/README.md) — the default: prices every path by the probability that all of its assumptions hold, and yields candidates lazily in exactly that order. diff --git a/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel new file mode 100644 index 00000000..7962c5c8 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel @@ -0,0 +1,29 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = [ + "bestfirst.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..50e2f99e --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/README.md @@ -0,0 +1,21 @@ +# Best-First Generator + +The default speculation path Generator. Every candidate path is a bet: building it pays off only if each of its dependency assumptions matches how that dependency actually resolves. This implementation prices each path by the probability of exactly that, and yields candidates in strictly non-increasing price order. + +## Pricing + +A path's price is the product, over its dependencies, of the probability that the assumed outcome happens. Each dependency's success probability comes from its state in the snapshot: `Succeeded` is a fact priced 1 and `Failed`/`Cancelled` are facts priced 0, so the contradicting branch is never generated; `Merging` (build passed, hand-off in flight) and `Cancelling` (halted, resurrected only by a lost cancel race) are not facts but are priced as modal certainties 1 and 0, suppressing the near-worthless opposite branch — if the long shot lands anyway, the controller refutes the affected paths and the next run re-plans against the new facts; `Created` and `Speculating` dependencies ask the injected scorer, once per batch per run. Scorer probabilities of exactly 0 or 1 pin the assumption the same way facts do. + +Two properties fall out of pure price ordering. Sure work streams first: a head whose dependencies are all decided has exactly one coherent path, priced 1, so real builds always outrank bets. And hedging is automatic: a dependency near probability 0.5 puts both of its branches adjacent in the stream, so a consumer with budget covers the head against either outcome. + +## Lazy enumeration + +Per head, the modal path takes every undecided dependency at its more likely outcome; every other path is that template with some subset of dependencies flipped, its price scaled by the flip ratios. Subsets are enumerated add/shift-style over flips sorted by descending ratio — each subset generated exactly once, each child priced at or below its parent — and a single max-heap merges all heads. The result is an exact global best-first stream produced in O(log n) per pull, materializing only what is pulled. + +Ties are deterministic: equal prices break by fewer flips (closer to the modal path), then by insertion order, and heads are seeded in queue order. + +## Coherence + +Dependencies are normalized into canonical queue order — ascending batch counter — before paths are built, because `entity.SpeculationPath` hashes the ordered dependencies into the path ID, and IDs must come out identical run after run for the queue to recognize paths it already built. Duplicate dependency entries collapse. A head whose snapshot is incomplete — a dependency missing, itself listed as its own dependency, or a dependency in a state the model does not cover — yields nothing this run and is re-planned when a later run sees a complete snapshot. + +See the [Best-First Speculation Generator RFC](../../../../../doc/rfc/submitqueue/speculation-generator.md) for the model, a worked example, and the design trade-offs. diff --git a/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go new file mode 100644 index 00000000..f505793a --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go @@ -0,0 +1,230 @@ +// 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 implements the default speculation path Generator. Every +// candidate path is a bet: building it pays off only if every one of its +// dependency assumptions matches how the dependency actually resolves. The +// generator prices each path by the probability of exactly that — resolved +// dependencies are pinned to their outcome, undecided ones get a success +// probability from the injected scorer — and yields candidates in strictly +// non-increasing price order, lazily, so the exponential path space is never +// materialized. +package bestfirst + +import ( + "cmp" + "context" + "fmt" + "math" + "slices" + "strconv" + "strings" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/scorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +// New returns the best-first Generator. The scorer supplies the success +// probability of undecided dependency batches; it is called at Open, at most +// once per batch per run, so anything expensive belongs behind the scorer's own +// cache. Wiring is trusted: a nil scorer panics on first use. +func New(sc scorer.Scorer) generator.Generator { + return &bestFirst{scorer: sc} +} + +// bestFirst builds one iterator per Open call; the generator itself carries no +// cross-run state, matching the speculate controller's recompute-from-scratch +// model. +type bestFirst struct { + // scorer supplies success probabilities for undecided dependency batches. + scorer scorer.Scorer +} + +// Open compiles every Speculating head in the snapshot into its lazy +// enumeration state and seeds the shared best-first heap with each head's most +// likely path. All scorer calls happen here; pulling from the iterator does no +// further I/O. +func (g *bestFirst) Open(ctx context.Context, batches []entity.Batch) (generator.PathIterator, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + byID := make(map[string]entity.Batch, len(batches)) + for _, b := range batches { + byID[b.ID] = b + } + + // Heads are seeded in queue order so that equally priced candidates — most + // notably the probability-1 paths of heads whose dependencies are all + // decided — stream oldest first. + heads := make([]entity.Batch, 0, len(batches)) + for _, b := range batches { + if b.State == entity.BatchStateSpeculating { + heads = append(heads, b) + } + } + slices.SortFunc(heads, func(a, b entity.Batch) int { return compareQueueOrder(a.ID, b.ID) }) + + it := &iterator{} + scores := make(map[string]float64, len(batches)) + for _, head := range heads { + hs, ok, err := g.prepareHead(ctx, head, byID, scores) + if err != nil { + return nil, err + } + if !ok { + continue + } + it.addHead(hs) + } + return it, nil +} + +// prepareHead compiles one head into its stream state: the modal path template, +// the swing dependencies in flip order, and the modal path's probability. +// +// ok is false when the snapshot cannot support coherent candidates for this +// head — a dependency is missing from the snapshot, is the head itself, or is +// in a state the model does not cover. The head then yields nothing this run; +// a later run sees a complete snapshot and plans it again. +func (g *bestFirst) prepareHead(ctx context.Context, head entity.Batch, byID map[string]entity.Batch, scores map[string]float64) (headStream, bool, error) { + // Dependency order must be canonical: entity.SpeculationPath hashes the + // ordered dependencies into the path ID, and a run that ordered them + // differently would re-propose paths the queue already built. The batch's + // own Dependencies order is unspecified, so it is normalized here. + depIDs := slices.Clone(head.Dependencies) + slices.SortFunc(depIDs, compareQueueOrder) + depIDs = slices.Compact(depIDs) + + hs := headStream{ + template: entity.SpeculationPath{ + Head: head.ID, + Dependencies: make([]entity.PathDependency, 0, len(depIDs)), + }, + baseProbability: 1, + } + for _, depID := range depIDs { + if depID == head.ID { + return headStream{}, false, nil + } + dep, present := byID[depID] + if !present { + return headStream{}, false, nil + } + p, known, err := g.successProbability(ctx, dep, scores) + if err != nil { + return headStream{}, false, err + } + if !known { + return headStream{}, false, nil + } + + modal, flipped := entity.DependencyAssumptionSucceeds, entity.DependencyAssumptionFails + if p < 0.5 { + modal, flipped = flipped, modal + } + hs.template.Dependencies = append(hs.template.Dependencies, entity.PathDependency{Batch: depID, Assumption: modal}) + + // Only genuinely undecided dependencies swing: a probability of exactly + // 0 or 1 pins the assumption, contributes no probability factor, and + // suppresses the opposite branch outright — its paths would be priced 0. + if p > 0 && p < 1 { + q := math.Max(p, 1-p) + hs.flips = append(hs.flips, flip{ + depIndex: len(hs.template.Dependencies) - 1, + assumption: flipped, + ratio: (1 - q) / q, + }) + hs.baseProbability *= q + } + } + // Cheapest flips first; the stable sort keeps queue order between equal + // ratios, so the whole stream stays deterministic. + slices.SortStableFunc(hs.flips, func(a, b flip) int { return cmp.Compare(b.ratio, a.ratio) }) + return hs, true, nil +} + +// successProbability is the probability the model assigns to dep resolving as +// succeeded. +// +// Resolved states are facts: 1 or 0, per the contract that a candidate never +// contradicts a resolved outcome. Merging and Cancelling are not facts, but +// their modal outcomes are lopsided enough to price as certainties — a Merging +// batch has already passed a build and been handed off to land, and a +// Cancelling batch is halted, resurrected only by a lost cancel race. Pricing +// them at 1 and 0 suppresses the near-worthless opposite branch; if the long +// shot lands anyway, the controller refutes the affected paths and the next +// run re-plans against the new facts, like any other lost bet. +// +// Undecided pipeline states ask the scorer, memoized per run. known is false +// for states the model does not cover (Creating, unknown), which are not +// eligible dependencies in the first place. +func (g *bestFirst) successProbability(ctx context.Context, dep entity.Batch, scores map[string]float64) (p float64, known bool, err error) { + switch dep.State { + case entity.BatchStateSucceeded, entity.BatchStateMerging: + return 1, true, nil + case entity.BatchStateFailed, entity.BatchStateCancelled, entity.BatchStateCancelling: + return 0, true, nil + case entity.BatchStateCreated, entity.BatchStateSpeculating: + if p, ok := scores[dep.ID]; ok { + return p, true, nil + } + p, err := g.scorer.Score(ctx, dep) + if err != nil { + return 0, false, fmt.Errorf("score dependency %s: %w", dep.ID, err) + } + if math.IsNaN(p) || p < 0 || p > 1 { + return 0, false, fmt.Errorf("scorer returned probability %v for batch %s, want a value in [0, 1]", p, dep.ID) + } + scores[dep.ID] = p + return p, true, nil + default: + return 0, false, nil + } +} + +// compareQueueOrder orders batch IDs by queue order: ascending counter for the +// documented "/batch/" ID format. Dependencies always share +// their head's queue, so the counter alone decides. IDs whose counter does not +// parse sort after all that do, then by plain string comparison, keeping the +// order total and deterministic for any input. +func compareQueueOrder(a, b string) int { + na, aok := batchCounter(a) + nb, bok := batchCounter(b) + switch { + case aok && bok: + if c := cmp.Compare(na, nb); c != 0 { + return c + } + return strings.Compare(a, b) + case aok: + return -1 + case bok: + return 1 + default: + return strings.Compare(a, b) + } +} + +// batchCounter extracts the numeric counter from a "/batch/" +// batch ID. +func batchCounter(id string) (int64, bool) { + i := strings.LastIndexByte(id, '/') + if i < 0 { + return 0, false + } + n, err := strconv.ParseInt(id[i+1:], 10, 64) + return n, err == 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..0dd36335 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go @@ -0,0 +1,396 @@ +// 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" + "errors" + "fmt" + "math" + "testing" + + "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" +) + +// scoreFunc adapts a function to scorer.Scorer. +type scoreFunc func(ctx context.Context, batch entity.Batch) (float64, error) + +func (f scoreFunc) Score(ctx context.Context, batch entity.Batch) (float64, error) { + return f(ctx, batch) +} + +// scoresByID serves fixed per-batch scores and fails the run on any batch it +// has no score for, so a test also asserts which batches get scored at all. +func scoresByID(scores map[string]float64) scorer.Scorer { + return scoreFunc(func(_ context.Context, b entity.Batch) (float64, error) { + s, ok := scores[b.ID] + if !ok { + return 0, fmt.Errorf("unexpected score request for %s", b.ID) + } + return s, nil + }) +} + +func batch(id string, state entity.BatchState, deps ...string) entity.Batch { + return entity.Batch{ID: id, Queue: "q", State: state, Dependencies: deps} +} + +func succeeds(id string) entity.PathDependency { + return entity.PathDependency{Batch: id, Assumption: entity.DependencyAssumptionSucceeds} +} + +func fails(id string) entity.PathDependency { + return entity.PathDependency{Batch: id, Assumption: entity.DependencyAssumptionFails} +} + +func path(head string, deps ...entity.PathDependency) entity.SpeculationPath { + return entity.SpeculationPath{Head: head, Dependencies: append([]entity.PathDependency{}, deps...)} +} + +// drain pulls until the stream ends, requiring it to end before max pulls. +func drain(t *testing.T, it generator.PathIterator, max int) []entity.CandidatePath { + t.Helper() + var out []entity.CandidatePath + for { + c, ok, err := it.Next(context.Background()) + require.NoError(t, err) + if !ok { + return out + } + out = append(out, c) + require.Less(t, len(out), max, "stream did not end") + } +} + +// requireStream asserts the exact candidate sequence: paths equal, scores close. +func requireStream(t *testing.T, want []entity.CandidatePath, got []entity.CandidatePath) { + t.Helper() + require.Len(t, got, len(want)) + for i := range want { + assert.Equalf(t, want[i].Path, got[i].Path, "candidate %d", i) + assert.InDeltaf(t, want[i].RankingScore, got[i].RankingScore, 1e-9, "candidate %d score", i) + } +} + +func TestSingleHeadStreams(t *testing.T) { + tests := []struct { + name string + batches []entity.Batch + scores map[string]float64 + want []entity.CandidatePath + }{ + { + name: "no dependencies yields the sure path", + batches: []entity.Batch{batch("q/batch/1", entity.BatchStateSpeculating)}, + want: []entity.CandidatePath{ + {Path: path("q/batch/1"), RankingScore: 1}, + }, + }, + { + name: "resolved dependencies are pinned facts and only undecided ones swing", + batches: []entity.Batch{ + batch("q/batch/1", entity.BatchStateSucceeded), + batch("q/batch/2", entity.BatchStateFailed), + batch("q/batch/3", entity.BatchStateCancelled), + batch("q/batch/4", entity.BatchStateCreated), + batch("q/batch/9", entity.BatchStateSpeculating, "q/batch/1", "q/batch/2", "q/batch/3", "q/batch/4"), + }, + scores: map[string]float64{"q/batch/4": 0.6}, + want: []entity.CandidatePath{ + {Path: path("q/batch/9", succeeds("q/batch/1"), fails("q/batch/2"), fails("q/batch/3"), succeeds("q/batch/4")), RankingScore: 0.6}, + {Path: path("q/batch/9", succeeds("q/batch/1"), fails("q/batch/2"), fails("q/batch/3"), fails("q/batch/4")), RankingScore: 0.4}, + }, + }, + { + name: "merging and cancelling dependencies are priced as modal certainties", + batches: []entity.Batch{ + batch("q/batch/1", entity.BatchStateMerging), + batch("q/batch/2", entity.BatchStateCancelling), + batch("q/batch/9", entity.BatchStateSpeculating, "q/batch/1", "q/batch/2"), + }, + want: []entity.CandidatePath{ + {Path: path("q/batch/9", succeeds("q/batch/1"), fails("q/batch/2")), RankingScore: 1}, + }, + }, + { + name: "scores of exactly zero and one pin the assumption", + batches: []entity.Batch{ + batch("q/batch/1", entity.BatchStateCreated), + batch("q/batch/2", entity.BatchStateCreated), + batch("q/batch/9", entity.BatchStateSpeculating, "q/batch/1", "q/batch/2"), + }, + scores: map[string]float64{"q/batch/1": 0, "q/batch/2": 1}, + want: []entity.CandidatePath{ + {Path: path("q/batch/9", fails("q/batch/1"), succeeds("q/batch/2")), RankingScore: 1}, + }, + }, + { + name: "coin-flip dependency hedges both branches, modal first", + batches: []entity.Batch{ + batch("q/batch/1", entity.BatchStateSpeculating), + batch("q/batch/9", entity.BatchStateSpeculating, "q/batch/1"), + }, + scores: map[string]float64{"q/batch/1": 0.5}, + want: []entity.CandidatePath{ + {Path: path("q/batch/1"), RankingScore: 1}, + {Path: path("q/batch/9", succeeds("q/batch/1")), RankingScore: 0.5}, + {Path: path("q/batch/9", fails("q/batch/1")), RankingScore: 0.5}, + }, + }, + { + name: "duplicate dependency entries collapse to one", + batches: []entity.Batch{ + batch("q/batch/1", entity.BatchStateCreated), + batch("q/batch/9", entity.BatchStateSpeculating, "q/batch/1", "q/batch/1"), + }, + scores: map[string]float64{"q/batch/1": 0.7}, + want: []entity.CandidatePath{ + {Path: path("q/batch/9", succeeds("q/batch/1")), RankingScore: 0.7}, + {Path: path("q/batch/9", fails("q/batch/1")), RankingScore: 0.3}, + }, + }, + { + name: "unlikely dependency is assumed to fail on the modal path", + batches: []entity.Batch{ + batch("q/batch/1", entity.BatchStateCreated), + batch("q/batch/9", entity.BatchStateSpeculating, "q/batch/1"), + }, + scores: map[string]float64{"q/batch/1": 0.2}, + want: []entity.CandidatePath{ + {Path: path("q/batch/9", fails("q/batch/1")), RankingScore: 0.8}, + {Path: path("q/batch/9", succeeds("q/batch/1")), RankingScore: 0.2}, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + it, err := New(scoresByID(tc.scores)).Open(context.Background(), tc.batches) + require.NoError(t, err) + requireStream(t, tc.want, drain(t, it, 100)) + }) + } +} + +// TestBestFirstAcrossHeads walks the reference example from the design doc: a +// five-batch queue (two roots, a diamond, a leaf) whose full stream must come +// out in exact best-first order — sure paths first, hedges adjacent to their +// modal twins, deeper long shots later. +func TestBestFirstAcrossHeads(t *testing.T) { + batches := []entity.Batch{ + batch("q/batch/1", entity.BatchStateSpeculating), + batch("q/batch/2", entity.BatchStateSpeculating), + batch("q/batch/3", entity.BatchStateSpeculating, "q/batch/1", "q/batch/2"), + batch("q/batch/4", entity.BatchStateSpeculating, "q/batch/2"), + batch("q/batch/5", entity.BatchStateSpeculating, "q/batch/1", "q/batch/2", "q/batch/3", "q/batch/4"), + } + scores := map[string]float64{ + "q/batch/1": 0.9, + "q/batch/2": 0.5, + "q/batch/3": 0.8, + "q/batch/4": 0.7, + } + + it, err := New(scoresByID(scores)).Open(context.Background(), batches) + require.NoError(t, err) + got := drain(t, it, 100) + + wantFirst := []entity.CandidatePath{ + {Path: path("q/batch/1"), RankingScore: 1}, + {Path: path("q/batch/2"), RankingScore: 1}, + {Path: path("q/batch/4", succeeds("q/batch/2")), RankingScore: 0.5}, + {Path: path("q/batch/4", fails("q/batch/2")), RankingScore: 0.5}, + {Path: path("q/batch/3", succeeds("q/batch/1"), succeeds("q/batch/2")), RankingScore: 0.45}, + {Path: path("q/batch/3", succeeds("q/batch/1"), fails("q/batch/2")), RankingScore: 0.45}, + {Path: path("q/batch/5", succeeds("q/batch/1"), succeeds("q/batch/2"), succeeds("q/batch/3"), succeeds("q/batch/4")), RankingScore: 0.252}, + {Path: path("q/batch/5", succeeds("q/batch/1"), fails("q/batch/2"), succeeds("q/batch/3"), succeeds("q/batch/4")), RankingScore: 0.252}, + {Path: path("q/batch/5", succeeds("q/batch/1"), succeeds("q/batch/2"), succeeds("q/batch/3"), fails("q/batch/4")), RankingScore: 0.108}, + {Path: path("q/batch/5", succeeds("q/batch/1"), fails("q/batch/2"), succeeds("q/batch/3"), fails("q/batch/4")), RankingScore: 0.108}, + } + require.GreaterOrEqual(t, len(got), len(wantFirst)) + requireStream(t, wantFirst, got[:len(wantFirst)]) + + // The whole space: 1 + 1 + 4 + 2 + 16 paths, each exactly once. + assert.Len(t, got, 24) + seen := make(map[string]bool, len(got)) + for _, c := range got { + id := c.Path.ID() + assert.False(t, seen[id], "path repeated: %+v", c.Path) + seen[id] = true + } + + // Best-first means globally non-increasing prices, all of them positive. + for i := 1; i < len(got); i++ { + assert.LessOrEqual(t, got[i].RankingScore, got[i-1].RankingScore, "candidate %d outranks its predecessor", i) + } + assert.Greater(t, got[len(got)-1].RankingScore, 0.0) +} + +func TestOnlySpeculatingBatchesAreHeads(t *testing.T) { + batches := []entity.Batch{ + batch("q/batch/1", entity.BatchStateCreated), + batch("q/batch/2", entity.BatchStateMerging), + batch("q/batch/3", entity.BatchStateSucceeded), + batch("q/batch/4", entity.BatchStateCancelling), + batch("q/batch/5", entity.BatchStateSpeculating), + } + it, err := New(scoresByID(nil)).Open(context.Background(), batches) + require.NoError(t, err) + requireStream(t, []entity.CandidatePath{{Path: path("q/batch/5"), RankingScore: 1}}, drain(t, it, 10)) +} + +func TestIncompleteSnapshotSkipsTheHead(t *testing.T) { + tests := []struct { + name string + batches []entity.Batch + scores map[string]float64 + want []entity.CandidatePath + }{ + { + name: "dependency missing from the snapshot", + batches: []entity.Batch{ + batch("q/batch/2", entity.BatchStateSpeculating, "q/batch/1"), + batch("q/batch/3", entity.BatchStateSpeculating, "q/batch/2"), + }, + scores: map[string]float64{"q/batch/2": 0.7}, + want: []entity.CandidatePath{ + {Path: path("q/batch/3", succeeds("q/batch/2")), RankingScore: 0.7}, + {Path: path("q/batch/3", fails("q/batch/2")), RankingScore: 0.3}, + }, + }, + { + name: "dependency in a state the model does not cover", + batches: []entity.Batch{ + batch("q/batch/1", entity.BatchStateCreating), + batch("q/batch/2", entity.BatchStateSpeculating, "q/batch/1"), + batch("q/batch/3", entity.BatchStateSpeculating), + }, + want: []entity.CandidatePath{ + {Path: path("q/batch/3"), RankingScore: 1}, + }, + }, + { + name: "head listed as its own dependency", + batches: []entity.Batch{ + batch("q/batch/1", entity.BatchStateSpeculating, "q/batch/1"), + batch("q/batch/2", entity.BatchStateSpeculating), + }, + want: []entity.CandidatePath{ + {Path: path("q/batch/2"), RankingScore: 1}, + }, + }, + } + for _, tc := range tests { + t.Run(tc.name, func(t *testing.T) { + it, err := New(scoresByID(tc.scores)).Open(context.Background(), tc.batches) + require.NoError(t, err) + requireStream(t, tc.want, drain(t, it, 10)) + }) + } +} + +func TestScorerFailuresFailOpen(t *testing.T) { + head := batch("q/batch/2", entity.BatchStateSpeculating, "q/batch/1") + dep := batch("q/batch/1", entity.BatchStateCreated) + + t.Run("scorer error propagates", func(t *testing.T) { + sentinel := errors.New("scorer down") + sc := scoreFunc(func(context.Context, entity.Batch) (float64, error) { return 0, sentinel }) + _, err := New(sc).Open(context.Background(), []entity.Batch{dep, head}) + require.ErrorIs(t, err, sentinel) + }) + + for _, invalid := range []float64{math.NaN(), -0.1, 1.5} { + t.Run(fmt.Sprintf("invalid probability %v", invalid), func(t *testing.T) { + sc := scoreFunc(func(context.Context, entity.Batch) (float64, error) { return invalid, nil }) + _, err := New(sc).Open(context.Background(), []entity.Batch{dep, head}) + require.Error(t, err) + }) + } +} + +func TestContextCancellation(t *testing.T) { + batches := []entity.Batch{batch("q/batch/1", entity.BatchStateSpeculating)} + cancelled, cancel := context.WithCancel(context.Background()) + cancel() + + t.Run("open aborts", func(t *testing.T) { + _, err := New(scoresByID(nil)).Open(cancelled, batches) + require.ErrorIs(t, err, context.Canceled) + }) + + t.Run("next ends the stream", func(t *testing.T) { + it, err := New(scoresByID(nil)).Open(context.Background(), batches) + require.NoError(t, err) + _, ok, err := it.Next(cancelled) + require.ErrorIs(t, err, context.Canceled) + assert.False(t, ok) + }) +} + +// TestCanonicalOrderAndDeterminism feeds the same queue twice with every input +// slice reversed. Both streams must be identical — dependencies normalized into +// queue order (numeric counter order, so batch/2 before batch/10), heads seeded +// in queue order — because path IDs hash the dependency order and must come out +// the same on every run. +func TestCanonicalOrderAndDeterminism(t *testing.T) { + scores := map[string]float64{"q/batch/2": 0.9, "q/batch/10": 0.6} + build := func(depOrder, batchOrder bool) []entity.Batch { + deps := []string{"q/batch/10", "q/batch/2"} + if depOrder { + deps = []string{"q/batch/2", "q/batch/10"} + } + batches := []entity.Batch{ + batch("q/batch/1", entity.BatchStateSpeculating), + batch("q/batch/2", entity.BatchStateSpeculating), + batch("q/batch/10", entity.BatchStateSpeculating), + batch("q/batch/11", entity.BatchStateSpeculating, deps...), + } + if batchOrder { + for i, j := 0, len(batches)-1; i < j; i, j = i+1, j-1 { + batches[i], batches[j] = batches[j], batches[i] + } + } + return batches + } + + first := build(false, false) + inputDeps := first[3].Dependencies + it, err := New(scoresByID(scores)).Open(context.Background(), first) + require.NoError(t, err) + got := drain(t, it, 20) + + want := []entity.CandidatePath{ + {Path: path("q/batch/1"), RankingScore: 1}, + {Path: path("q/batch/2"), RankingScore: 1}, + {Path: path("q/batch/10"), RankingScore: 1}, + {Path: path("q/batch/11", succeeds("q/batch/2"), succeeds("q/batch/10")), RankingScore: 0.54}, + {Path: path("q/batch/11", succeeds("q/batch/2"), fails("q/batch/10")), RankingScore: 0.36}, + {Path: path("q/batch/11", fails("q/batch/2"), succeeds("q/batch/10")), RankingScore: 0.06}, + {Path: path("q/batch/11", fails("q/batch/2"), fails("q/batch/10")), RankingScore: 0.04}, + } + requireStream(t, want, got) + + // The caller's snapshot is read, never written: the head's dependency slice + // keeps its original order. + assert.Equal(t, []string{"q/batch/10", "q/batch/2"}, inputDeps) + + it2, err := New(scoresByID(scores)).Open(context.Background(), build(true, true)) + require.NoError(t, err) + require.Equal(t, got, drain(t, it2, 20)) +} diff --git a/submitqueue/extension/speculation/generator/bestfirst/iterator.go b/submitqueue/extension/speculation/generator/bestfirst/iterator.go new file mode 100644 index 00000000..a8a970f4 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/iterator.go @@ -0,0 +1,198 @@ +// 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" +) + +// headStream is one head's compiled enumeration state. Every path of the head +// is the modal template with some subset of the flips applied, priced at +// baseProbability times the product of the applied flip ratios — which is what +// lets the iterator walk the 2^k subsets best-first without materializing them. +type headStream struct { + // template is the head's modal path: every dependency in queue order at its + // most likely assumption. + template entity.SpeculationPath + // flips are the head's swing dependencies — those whose outcome is genuinely + // undecided — in enumeration order: descending ratio, queue order on ties. + flips []flip + // baseProbability prices the modal path: the product of every swing + // dependency's modal-outcome probability. Pinned dependencies are certain + // and contribute no factor. + baseProbability float64 +} + +// flip is one swing dependency's less likely branch. +type flip struct { + // depIndex locates the dependency inside template.Dependencies. + depIndex int + // assumption is the flipped, less likely assumption. + assumption entity.DependencyAssumption + // ratio multiplies a path's price when this flip is applied: (1−q)/q for + // modal-outcome probability q. It is in (0, 1] because swing dependencies + // have probability strictly between 0 and 1. + ratio float64 +} + +// flipNode is one link of an entry's immutable flip chain. Entries share their +// chain tails with the parent they were derived from, so an entry costs O(1) +// memory regardless of how many flips it carries. +type flipNode struct { + // idx indexes the head's flips. + idx int + // prev is the rest of the chain, nil at the start. + prev *flipNode +} + +// entry is one node of a head's enumeration tree: a flip subset represented by +// its chain, extended add/shift-style through lastFlip. +type entry struct { + // headIdx indexes iterator.heads. + headIdx int + // flips is the applied flip subset, nil for the modal path. + flips *flipNode + // lastFlip is the largest applied flip index, -1 for the modal path. + // Children either add lastFlip+1 to the subset or shift lastFlip up to it, + // which generates every subset exactly once. + lastFlip int + // nFlips is the subset size, the first tie-break: between equally priced + // candidates, the one closer to its modal path streams first. + nFlips int + // probability prices the materialized path — the chance that every + // assumption on it holds — and doubles as the candidate's ranking score. + probability float64 + // seq is the entry's insertion sequence, the final tie-break; it makes the + // stream fully deterministic and, all else equal, favors entries seeded or + // derived earlier. + seq int +} + +// iterator merges every head's lazy enumeration through one max-heap: pop the +// best entry, materialize it, push its at most two children. A child's +// probability never exceeds its parent's, so popped prices are non-increasing +// and the stream is exactly best-first. +type iterator struct { + // heads holds the per-head enumeration state entries index into. + heads []headStream + // entries is the frontier: for every head, the best not-yet-yielded subsets. + entries entryHeap + // seq numbers heap insertions to keep ties deterministic. + seq int +} + +// addHead seeds a head's stream with its modal path. +func (it *iterator) addHead(hs headStream) { + it.heads = append(it.heads, hs) + it.push(entry{ + headIdx: len(it.heads) - 1, + lastFlip: -1, + probability: hs.baseProbability, + }) +} + +func (it *iterator) push(e entry) { + e.seq = it.seq + it.seq++ + heap.Push(&it.entries, e) +} + +// Next yields the next candidate best-first. It never repeats a path: each +// heap entry is a distinct flip subset of its head, generated exactly once. +func (it *iterator) Next(ctx context.Context) (entity.CandidatePath, bool, error) { + if err := ctx.Err(); err != nil { + return entity.CandidatePath{}, false, err + } + if it.entries.Len() == 0 { + return entity.CandidatePath{}, false, nil + } + e := heap.Pop(&it.entries).(entry) + it.pushChildren(e) + return entity.CandidatePath{Path: it.materialize(e), RankingScore: e.probability}, true, nil +} + +// pushChildren pushes the popped entry's enumeration children: add extends the +// subset with the next flip, shift replaces the last flip with the next one. +// Flips are ordered by descending ratio, so both multiplications keep a child +// priced at or below its parent. +func (it *iterator) pushChildren(e entry) { + flips := it.heads[e.headIdx].flips + next := e.lastFlip + 1 + if next >= len(flips) { + return + } + it.push(entry{ + headIdx: e.headIdx, + flips: &flipNode{idx: next, prev: e.flips}, + lastFlip: next, + nFlips: e.nFlips + 1, + probability: e.probability * flips[next].ratio, + }) + if e.lastFlip < 0 { + return + } + it.push(entry{ + headIdx: e.headIdx, + flips: &flipNode{idx: next, prev: e.flips.prev}, + lastFlip: next, + nFlips: e.nFlips, + probability: e.probability / flips[e.lastFlip].ratio * flips[next].ratio, + }) +} + +// materialize renders an entry into a self-contained path: a copy of the modal +// template with the entry's flips applied. +func (it *iterator) materialize(e entry) entity.SpeculationPath { + hs := it.heads[e.headIdx] + deps := make([]entity.PathDependency, len(hs.template.Dependencies)) + copy(deps, hs.template.Dependencies) + for n := e.flips; n != nil; n = n.prev { + f := hs.flips[n.idx] + deps[f.depIndex].Assumption = f.assumption + } + return entity.SpeculationPath{Head: hs.template.Head, Dependencies: deps} +} + +// entryHeap orders entries highest probability first, then fewest flips, then +// lowest insertion sequence. +type entryHeap []entry + +func (h entryHeap) Len() int { return len(h) } + +func (h entryHeap) Less(i, j int) bool { + a, b := h[i], h[j] + if a.probability != b.probability { + return a.probability > b.probability + } + if a.nFlips != b.nFlips { + return a.nFlips < b.nFlips + } + return a.seq < b.seq +} + +func (h entryHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } + +func (h *entryHeap) Push(x any) { *h = append(*h, x.(entry)) } + +func (h *entryHeap) Pop() any { + old := *h + n := len(old) + x := old[n-1] + *h = old[:n-1] + return x +} diff --git a/submitqueue/extension/speculation/generator/generator.go b/submitqueue/extension/speculation/generator/generator.go new file mode 100644 index 00000000..8268571e --- /dev/null +++ b/submitqueue/extension/speculation/generator/generator.go @@ -0,0 +1,59 @@ +// 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 inside the +// standard Speculator. A Generator turns one snapshot of a queue's batches into +// a pull-based stream of candidate speculation paths, ranked by a policy of its +// choosing. It is not a controller-facing extension: the speculate controller +// depends only on the Speculator contract, and an alternate Speculator need not +// use or expose this seam. +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 +// one snapshot of the queue's batches. Which candidates come first, and what +// their ranking scores mean, is the implementation'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 or was cancelled 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) +} diff --git a/submitqueue/extension/storage/mock/request_batch_store_mock.go b/submitqueue/extension/storage/mock/request_batch_store_mock.go index 2c4c15f0..24e87e12 100644 --- a/submitqueue/extension/storage/mock/request_batch_store_mock.go +++ b/submitqueue/extension/storage/mock/request_batch_store_mock.go @@ -1,9 +1,9 @@ // Code generated by MockGen. DO NOT EDIT. -// Source: submitqueue/extension/storage/request_batch_store.go +// Source: request_batch_store.go // // Generated by this command: // -// mockgen -source=submitqueue/extension/storage/request_batch_store.go -destination=submitqueue/extension/storage/mock/request_batch_store_mock.go -package=mock +// mockgen -source=request_batch_store.go -destination=mock/request_batch_store_mock.go -package=mock // // Package mock is a generated GoMock package.