From 97047c435f485ca9263d49a61d1bfaec1758306f Mon Sep 17 00:00:00 2001 From: sergeyb Date: Wed, 5 Aug 2026 21:52:16 +0000 Subject: [PATCH] feat(speculation): add best-first path generator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: Intent: - Give the default speculator a lazy global stream of candidate paths ranked by the likelihood that dependency assumptions hold. - Avoid eager exponential path generation while preserving resolved dependency facts and deterministic ordering. Changes: - Add generator and iterator contracts plus a best-first implementation using log-space extend/swap enumeration and a global heap merge. - Add generated mocks, Bazel targets, comprehensive tests, and a focused design RFC with a worked example. --- Generated by the πŸͺ„ [pr-create](https://sg.uberinternal.com/code.uber.internal/uber-code/devexp-agent-marketplace/-/blob/claude-code/plugins/dev/uber-dev/skills/pr-create/SKILL.md) skill in devexp-agent-marketplace --- doc/rfc/index.md | 1 + doc/rfc/submitqueue/speculation-generator.md | 323 ++++++++++ doc/rfc/submitqueue/speculation.md | 2 +- .../speculation/generator/BUILD.bazel | 9 + .../extension/speculation/generator/README.md | 9 + .../generator/bestfirst/BUILD.bazel | 24 + .../speculation/generator/bestfirst/README.md | 9 + .../generator/bestfirst/bestfirst.go | 461 ++++++++++++++ .../generator/bestfirst/bestfirst_test.go | 571 ++++++++++++++++++ .../speculation/generator/generator.go | 47 ++ .../speculation/generator/mock/BUILD.bazel | 13 + .../generator/mock/generator_mock.go | 98 +++ 12 files changed, 1566 insertions(+), 1 deletion(-) 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/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..2fa9300a 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 +- [Probability-Ordered Speculation Generator](submitqueue/speculation-generator.md) - Lazy best-first enumeration of dependency assignments and global merging across Speculating heads - [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..21507c83 --- /dev/null +++ b/doc/rfc/submitqueue/speculation-generator.md @@ -0,0 +1,323 @@ +# Probability-Ordered Speculation Generator + +## Status + +Implemented. + +## Summary + +SubmitQueue's default speculation implementation needs one global stream of candidate paths across every batch currently in `BatchStateSpeculating`. A candidate path chooses one assumption for every dependency of its head: the dependency succeeds or it fails. Building every combination eagerly is infeasible because a head with `N` unresolved dependencies has `2^N` paths. + +The `bestfirst` generator produces this stream lazily. It scores unresolved dependencies once, enumerates each head's assignments in descending probability without materializing the full power set, and merges those per-head streams into one global best-first iterator. + +The implementation lives under `submitqueue/extension/speculation/generator`: the parent package defines the generator and iterator contracts, while `generator/bestfirst` contains the default probability-ranked implementation. + +## Goals + +1. Yield coherent candidate paths only for heads in `BatchStateSpeculating`. +2. Never contradict a dependency whose outcome is already terminal. +3. Rank candidates by the probability that all unresolved dependency assumptions match their eventual outcomes. +4. Rank globally across disconnected components and multiple heads. +5. Generate only the prefix requested by the allocator rather than all `2^N` paths. +6. Keep ordering numerically stable for deep dependency sets. +7. Avoid scoring the same shared dependency more than once per speculation run. +8. Preserve the head's dependency order in every emitted path. + +## Non-goals + +- The generator does not spend the build budget; that belongs to the allocator. +- The generator does not inspect path sets or suppress previously materialized paths; the allocator reconciles candidates with stored path sets. +- The generator does not propose cancellations, merges, or failure verdicts. +- The default implementation does not model correlation between dependency outcomes. +- The default implementation does not emit `DependencyAssumptionIgnored`; conflict relaxation requires a separate policy. + +## Contract + +Generation receives one queue snapshot containing all in-flight batches plus terminal batches still referenced as dependencies. Only batches in `BatchStateSpeculating` become heads. All other batches are facts or probability inputs. + +The snapshot must include every batch named by a generated head's `Dependencies` field. Missing dependencies, duplicate batch IDs, repeated dependencies, self-dependencies, empty IDs, and dependencies in the zero-value unknown state are malformed input and abort generation. + +The returned iterator yields `entity.CandidatePath` values. Exhaustion is represented by `ok=false`, not an error. Generation and iteration both honor context cancellation, and a cancelled `Next` call does not consume a candidate. + +## Dependency outcomes + +Terminal dependency states are evidence rather than predictions: + +| Batch state | Path assumption | +|---|---| +| `BatchStateSucceeded` | `DependencyAssumptionSucceeds` | +| `BatchStateFailed` | `DependencyAssumptionFails` | +| `BatchStateCancelled` | `DependencyAssumptionFails` | + +Every other nonzero state remains unresolved and is scored. In particular, `BatchStateCancelling` remains unresolved because cancellation is best-effort and a merge may still win the race. + +Every dependency remains present in the emitted path, including resolved dependencies. This keeps the path self-describing while ensuring every candidate agrees with known facts. + +## Probability model + +For an unresolved dependency `d`, the injected scorer supplies: + +```text +p[d] = P(d eventually succeeds) +``` + +The default generator assumes unresolved outcomes are independent. For candidate `c`, let `U(c)` be its unresolved dependencies: + +```text +P(c) = product over d in U(c): + p[d] when c assumes d succeeds + 1 - p[d] when c assumes d fails +``` + +Resolved dependencies do not contribute a factor because generation is already conditioned on their known outcomes. + +The head's own score is not included. A path result is useful whether the head's validation passes or fails; the ranking question is whether the path's dependency assumptions become the actual world. + +Scorer outputs must be finite values in the inclusive range `[0, 1]`. Invalid values abort the run rather than silently changing the model. Exact `0` and `1` are preserved: paths betting against them receive ranking score `0`. + +## Architecture + +```text + Queue snapshot + β”‚ + β”œβ”€β”€ validate IDs, states, and dependency references + β”‚ + β”œβ”€β”€ score every shared unresolved dependency once + β”‚ + └── build one lazy stream per Speculating head + β”‚ + β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ Per-head assignment streamβ”‚ + β”‚ MAP path, then cheapest β”‚ + β”‚ combinations of flips β”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ one current candidate per head + β–Ό + β”Œβ”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β” + β”‚ Global max-heap β”‚ + β”‚ ordered by log probabilityβ”‚ + β””β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”¬β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”€β”˜ + β”‚ + β–Ό + Iterator.Next() +``` + +The global heap contains only the current best candidate from each head. When `Next` removes a candidate, it advances only that candidate's local stream and inserts the head's next candidate. This is a k-way merge over already ordered streams. + +## Per-head best-first enumeration + +### Optimal assignment + +For each unresolved dependency: + +```text +preferred(d) = succeeds when p[d] >= 0.5 + fails otherwise +``` + +Choosing every dependency's preferred outcome gives the maximum-probability assignment for the head. + +### Flips + +Every other assignment is the optimal assignment with some dependencies flipped to their less likely outcomes. + +For probabilities strictly between `0` and `1`, flipping dependency `d` has log penalty: + +```text +flipCost(d) = log(max(p[d], 1-p[d]) / min(p[d], 1-p[d])) +``` + +The probability of a flipped assignment is: + +```text +log P(assignment) = log P(optimal) - sum(flipCost(d) for flipped d) +``` + +Dependencies are sorted by ascending flip cost. A dependency near `0.5` is cheap to flip because the scheduler is uncertain about it; a dependency near `0` or `1` is expensive to flip. + +### Extend/swap subset enumeration + +The local min-heap stores subsets of sorted dependencies to flip. After popping a subset whose greatest sorted index is `j`, the generator creates at most two successors: + +```text +extend: keep the current flips and also flip j+1 +swap: replace flip j with flip j+1 +``` + +Pseudocode: + +```text +emit optimal assignment +seed heap with {0} + +while heap is not empty: + current = pop lowest flip cost + emit assignment produced by current.flips + + j = current.lastIndex + if j+1 exists: + push current.flips + {j+1} + push current.flips - {j} + {j+1} +``` + +The extend/swap tree reaches every subset exactly once. Because flip costs are sorted, no generated child has a better nonzero probability than its parent. + +Exact zero-probability outcomes require special handling because an infinite flip cost makes `infinity - infinity` undefined during swap. The implementation tracks zero-probability flips as an integer count alongside the finite log cost. Any positive count produces probability zero while preserving complete, deterministic enumeration. + +## Numerical ordering + +The iterator orders candidates using internal log probabilities, not the public `RankingScore`. + +Multiplying many probabilities can underflow to zero, and converting a very negative log probability back to `float64` can also underflow. Two deep paths may therefore both expose `RankingScore=0` even though one is mathematically more likely. The internal log values remain distinct and preserve the correct iteration order. + +Equal-probability candidates use deterministic tie-breaking. Heads are initialized by ID, local equal-cost entries use insertion sequence, and the global heap uses head and path identity. + +## Worked example + +Consider two disconnected components. Arrows point from a dependency to its dependent. + +```text +Component 1 + +A (0.80) ──┐ + β”œβ”€β”€β–Ά C +B (0.60) β”€β”€β”˜ + +Component 2 + +E (0.30) ─────▢ F +``` + +`C` and `F` are Speculating heads. `A`, `B`, and `E` are unresolved dependencies. The values in parentheses are scorer probabilities of success. + +### Step 1: create `C`'s local stream + +`C` has four assignments: + +| Rank | Assignment | Probability | +|---:|---|---:| +| 1 | `A=succeeds, B=succeeds` | `0.80 * 0.60 = 0.48` | +| 2 | `A=succeeds, B=fails` | `0.80 * 0.40 = 0.32` | +| 3 | `A=fails, B=succeeds` | `0.20 * 0.60 = 0.12` | +| 4 | `A=fails, B=fails` | `0.20 * 0.40 = 0.08` | + +The optimal assignment is `A=succeeds, B=succeeds`. Flipping `B` is cheaper than flipping `A`, so the second candidate changes only `B`. + +### Step 2: create `F`'s local stream + +`E` is more likely to fail than succeed: + +| Rank | Assignment | Probability | +|---:|---|---:| +| 1 | `E=fails` | `0.70` | +| 2 | `E=succeeds` | `0.30` | + +### Step 3: initialize the global heap + +The global heap initially contains only each head's first candidate: + +```text +F: E=fails 0.70 +C: A=succeeds, B=succeeds 0.48 +``` + +### Step 4: pull the first candidate + +`Next` returns `F: E=fails` with score `0.70`, advances only `F`, and inserts `F: E=succeeds` with score `0.30`. + +```text +C: A=succeeds, B=succeeds 0.48 +F: E=succeeds 0.30 +``` + +### Step 5: pull the second candidate + +`Next` returns `C: A=succeeds, B=succeeds` with score `0.48`, advances only `C`, and inserts `C: A=succeeds, B=fails` with score `0.32`. + +```text +C: A=succeeds, B=fails 0.32 +F: E=succeeds 0.30 +``` + +### Step 6: continue the merge + +The final global order is: + +```text +1. F: E=fails 0.70 +2. C: A=succeeds, B=succeeds 0.48 +3. C: A=succeeds, B=fails 0.32 +4. F: E=succeeds 0.30 +5. C: A=fails, B=succeeds 0.12 +6. C: A=fails, B=fails 0.08 +``` + +No disconnected-component special case is required; the global heap naturally merges them. + +## Evidence update example + +The generator is intentionally run-scoped rather than stateful across queue updates. Suppose a later speculation run sees `E` resolved to succeeded. + +`F` then has one coherent path: + +```text +F: E=succeeds score 1.0 +``` + +The previous `E=fails` candidate is no longer generated because it contradicts a resolved fact. The allocator and controller reconcile any previously materialized path against the new snapshot. + +## Correctness properties + +### Coherence + +Resolved dependency outcomes map to exactly one assumption, and unresolved dependencies map to exactly one of succeeds/fails in every candidate. Therefore each emitted path is complete and cannot contradict known evidence. + +### No duplicate assignments + +Every unresolved assignment corresponds to one subset of variables flipped from the optimal assignment. The extend/swap tree enumerates every subset exactly once, so a local stream cannot repeat a path. + +### Per-head ordering + +For nonzero paths, assignment probability decreases monotonically as total flip cost increases. The local min-heap therefore emits assignments in descending probability order. Zero-probability paths follow all nonzero paths. + +### Global ordering + +The global heap contains the next unconsumed item from every per-head ordered stream. Removing the greatest item and replacing it with the same stream's successor is the standard k-way merge invariant, so the returned sequence is globally ordered. + +## Complexity + +Let `D` be the number of unique unresolved dependency batches referenced by Speculating heads, `H` the number of Speculating heads, `N_h` the unresolved dependency count for head `h`, and `K` the number of candidates actually pulled. + +- Snapshot indexing and validation: `O(number of batches + dependency references)`. +- Scoring: `D` scorer calls. +- Per-head initialization: `O(N_h log N_h)` to sort flip costs. +- Global initialization: `O(H)`. +- Each pull: global heap work `O(log H)`, local heap work up to `O(log K_h)`, plus `O(number of head dependencies)` to materialize the self-describing output path. +- Enumeration state grows with the consumed prefix rather than `2^N`. + +The worst case remains exponential if a caller exhausts every path. The allocator is expected to pull only enough candidates to fill or compare against the finite build budget. + +## Alternatives considered + +### Eager power-set generation + +Generating all assignments and sorting them is simple but requires exponential time and memory before the first candidate can be consumed. + +### Probability threshold + +A minimum score can prune low-probability branches, but output size becomes distribution-dependent and the generator cannot guarantee that it has found the best finite prefix without exploring all branches above the threshold. + +### Persisted ranking + +Ranking scores depend on the current queue snapshot and become stale when any dependency resolves or receives a new score. Scores therefore remain transient on `CandidatePath` and are never persisted. + +### Include the head's score + +Multiplying by the head's own success probability would prioritize paths likely to produce a passing build rather than paths likely to match reality. The generator's contract ranks dependency assumptions, so it excludes the head score. An alternate generator may choose a different ranking policy without changing correctness. + +## Future work + +- A correlated-outcome generator can replace the independent product model while retaining the same iterator contract. +- A conflict-relaxing generator can add `DependencyAssumptionIgnored` choices and rank them using an explicit relaxation policy. +- A graph-impact ranking can combine assignment probability with critical-path or business value when the scheduling objective changes from expected applicability to full-queue makespan. diff --git a/doc/rfc/submitqueue/speculation.md b/doc/rfc/submitqueue/speculation.md index a0a744d9..0b97ed90 100644 --- a/doc/rfc/submitqueue/speculation.md +++ b/doc/rfc/submitqueue/speculation.md @@ -101,7 +101,7 @@ The one extension. It decides *which paths to build and which running ones to ca The default Speculator is composed from two swappable interfaces β€” a **Generator** and an **Allocator** β€” so scoring and preemption policy can vary independently. They are composition points inside the default implementation, not controller-facing extensions: the controller depends only on the Speculator contract, and an alternate Speculator need not use or expose this split. The default opens the Generator's candidate stream over the batches, then hands that stream and the path sets to the Allocator. -- **Generator** β€” yields the queue's candidate paths as one iterator across heads in `BatchStateSpeculating`. *Contract:* every candidate has a Speculating head and is coherent; none repeats or contradicts a resolved fact. Ranking is implementation-defined β€” the Generator may compute it directly, call an injected scorer extension, or use other injected data β€” and the score it carries is meaningful only within the run. The Allocator consumes the iterator in the order the Generator yields it and does not interpret the score. *Default:* `bestfirst` ranks best-first by the probability that a path's assumptions all hold. +- **Generator** β€” yields the queue's candidate paths as one iterator across heads in `BatchStateSpeculating`. *Contract:* every candidate has a Speculating head and is coherent; none repeats or contradicts a resolved fact. Ranking is implementation-defined β€” the Generator may compute it directly, call an injected scorer extension, or use other injected data β€” and the score it carries is meaningful only within the run. The Allocator consumes the iterator in the order the Generator yields it and does not interpret the score. *Default:* `bestfirst` ranks best-first by the probability that a path's assumptions all hold; its lazy enumeration and global merge are described in [Probability-Ordered Speculation Generator](speculation-generator.md). - **Allocator** β€” spends the build budget (the queue's cap on concurrent builds) over the iterator. *Contract:* it pulls in order until the budget fills and matches candidates to existing paths by ID, so a pending or building path keeps the slot it already holds rather than starting a second attempt, and a candidate whose path is already terminal in the path sets is skipped rather than rebuilt; pending dispatches are replayed by the controller as described above. Pending, building, and cancelling paths charge the budget (a cancelling build holds CI until terminal), while terminal ones charge none. Cancellation is best-effort, so the Allocator does not spend capacity it merely expects a cancel to release and risk exceeding the hard CI cap. *Default:* the sticky policy fills only free slots and leaves in-flight builds running; a preempting policy cancels in-flight paths below the funded set. Budget is the only rationing lever β€” there is no ranking-score floor. A build cancelled to make room still charges budget until its cancel reaches terminal and publishes dirty, so the next run funds the released slot β€” the queue converges over successive ticks rather than oversubscribing in a single pass. ### Extension APIs 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..c1d0eb95 --- /dev/null +++ b/submitqueue/extension/speculation/generator/README.md @@ -0,0 +1,9 @@ +# generator + +The `generator` package defines the pull-based candidate stream consumed by the default speculation allocator. A generator receives one queue snapshot and yields coherent paths for heads in `BatchStateSpeculating`, ordered according to implementation-owned ranking scores. + +Every yielded path contains exactly one assumption for every dependency of its head. Terminal dependency outcomes are fixed facts: succeeded dependencies are assumed to succeed, while failed and cancelled dependencies are assumed to fail. A generator never emits an assumption that contradicts those facts. + +The iterator reports exhaustion through its boolean result rather than an error. Both generation and iteration honor context cancellation. Iterators own mutable traversal state and are not safe for concurrent use. + +`bestfirst` is the default implementation. It treats unresolved dependency outcomes as independent, obtains their success probabilities from an injected scorer, and lazily merges per-head streams in descending assignment-probability 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..dd0d1410 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/BUILD.bazel @@ -0,0 +1,24 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["bestfirst.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/generator/bestfirst", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/scorer:go_default_library", + "//submitqueue/extension/speculation/generator:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = ["bestfirst_test.go"], + embed = [":go_default_library"], + deps = [ + "//submitqueue/entity:go_default_library", + "@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..2ec0f700 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/README.md @@ -0,0 +1,9 @@ +# bestfirst + +`bestfirst` is the default speculation candidate generator. It scores each unresolved dependency at most once per queue snapshot, treats those outcomes as independent, and ranks a path by the probability that all of its assumptions match the eventual dependency outcomes. + +Each head is represented by a lazy local stream. The stream starts with the most likely assignment and enumerates successively less likely assignments by applying combinations of outcome flips in log-probability order. A global heap merges the local streams so callers receive one best-first sequence across all Speculating heads and disconnected graph components. + +Resolved dependencies are fixed rather than scored. Succeeded dependencies always use the succeeds assumption; failed and cancelled dependencies always use the fails assumption. Cancelling dependencies remain probabilistic because cancellation is best-effort. + +See the probability-ordered speculation generator RFC for the algorithm, numerical treatment, worked example, and complexity analysis. diff --git a/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go new file mode 100644 index 00000000..80aeb5db --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst.go @@ -0,0 +1,461 @@ +// Copyright (c) 2026 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package bestfirst provides a probability-ordered speculation path generator. +package bestfirst + +import ( + "container/heap" + "context" + "fmt" + "math" + "sort" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/scorer" + "github.com/uber/submitqueue/submitqueue/extension/speculation/generator" +) + +// bestFirst generates candidate paths using independent dependency +// probabilities supplied by scorer. +type bestFirst struct { + scorer scorer.Scorer +} + +var _ generator.Generator = (*bestFirst)(nil) + +// New returns a Generator that ranks paths by the probability that every +// unresolved dependency assumption holds. The scorer is called at most once +// per unresolved dependency batch in each Generate call. +func New(s scorer.Scorer) generator.Generator { + if s == nil { + panic("bestfirst.New: scorer must not be nil") + } + return &bestFirst{scorer: s} +} + +// Generate validates the queue snapshot, resolves the dependency probabilities +// needed by Speculating heads, and opens a lazy global best-first iterator. +func (g *bestFirst) Generate(ctx context.Context, batches []entity.Batch) (generator.Iterator, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + batchByID := make(map[string]entity.Batch, len(batches)) + for _, batch := range batches { + if batch.ID == "" { + return nil, fmt.Errorf("bestfirst: batch has empty ID") + } + if _, exists := batchByID[batch.ID]; exists { + return nil, fmt.Errorf("bestfirst: duplicate batch ID %q", batch.ID) + } + batchByID[batch.ID] = batch + } + + heads := make([]entity.Batch, 0) + unresolvedDependencyIDs := make(map[string]struct{}) + for _, batch := range batches { + if batch.State != entity.BatchStateSpeculating { + continue + } + heads = append(heads, batch) + + seen := make(map[string]struct{}, len(batch.Dependencies)) + for _, dependencyID := range batch.Dependencies { + if dependencyID == "" { + return nil, fmt.Errorf("bestfirst: head %q has an empty dependency ID", batch.ID) + } + if dependencyID == batch.ID { + return nil, fmt.Errorf("bestfirst: head %q depends on itself", batch.ID) + } + if _, duplicate := seen[dependencyID]; duplicate { + return nil, fmt.Errorf("bestfirst: head %q repeats dependency %q", batch.ID, dependencyID) + } + seen[dependencyID] = struct{}{} + + dependency, exists := batchByID[dependencyID] + if !exists { + return nil, fmt.Errorf("bestfirst: head %q references missing dependency %q", batch.ID, dependencyID) + } + if dependency.State == entity.BatchStateUnknown { + return nil, fmt.Errorf("bestfirst: dependency %q has unknown state", dependencyID) + } + if _, resolved := resolvedAssumption(dependency.State); !resolved { + unresolvedDependencyIDs[dependencyID] = struct{}{} + } + } + } + sort.Slice(heads, func(i, j int) bool { + return heads[i].ID < heads[j].ID + }) + + probabilityByID := make(map[string]float64, len(unresolvedDependencyIDs)) + ids := make([]string, 0, len(unresolvedDependencyIDs)) + for id := range unresolvedDependencyIDs { + ids = append(ids, id) + } + sort.Strings(ids) + for _, id := range ids { + if err := ctx.Err(); err != nil { + return nil, err + } + probability, err := g.scorer.Score(ctx, batchByID[id]) + if err != nil { + return nil, fmt.Errorf("bestfirst: score dependency %q: %w", id, err) + } + if math.IsNaN(probability) || math.IsInf(probability, 0) || probability < 0 || probability > 1 { + return nil, fmt.Errorf("bestfirst: dependency %q has invalid probability %v", id, probability) + } + probabilityByID[id] = probability + } + + it := &candidateIterator{} + heap.Init(&it.candidates) + for _, head := range heads { + stream := newPathStream(head, batchByID, probabilityByID) + candidate, ok := stream.next() + if !ok { + continue + } + heap.Push(&it.candidates, candidateItem{ + rankedCandidate: candidate, + pathID: candidate.candidate.Path.ID(), + stream: stream, + }) + } + return it, nil +} + +// resolvedAssumption converts a terminal dependency outcome into the only +// coherent path assumption. Cancelling remains unresolved because cancellation +// is best-effort and the batch may still succeed. +func resolvedAssumption(state entity.BatchState) (entity.DependencyAssumption, bool) { + switch state { + case entity.BatchStateSucceeded: + return entity.DependencyAssumptionSucceeds, true + case entity.BatchStateFailed, entity.BatchStateCancelled: + return entity.DependencyAssumptionFails, true + default: + return entity.DependencyAssumptionUnknown, false + } +} + +// dependencyVariable is one unresolved dependency ordered by the penalty for +// flipping away from its most likely outcome. +type dependencyVariable struct { + dependencyIndex int + flipCost float64 + zeroProbability bool +} + +// pathStream lazily enumerates one head's paths in descending probability. +type pathStream struct { + head string + base []entity.PathDependency + variables []dependencyVariable + optimalLogProbability float64 + emitOptimal bool + subsets flipHeap + nextSequence uint64 +} + +func newPathStream( + head entity.Batch, + batchByID map[string]entity.Batch, + probabilityByID map[string]float64, +) *pathStream { + stream := &pathStream{ + head: head.ID, + base: make([]entity.PathDependency, len(head.Dependencies)), + emitOptimal: true, + } + + for i, dependencyID := range head.Dependencies { + dependency := batchByID[dependencyID] + if assumption, resolved := resolvedAssumption(dependency.State); resolved { + stream.base[i] = entity.PathDependency{ + Batch: dependencyID, + Assumption: assumption, + } + continue + } + + probability := probabilityByID[dependencyID] + preferredProbability := math.Max(probability, 1-probability) + nonPreferredProbability := math.Min(probability, 1-probability) + preferredAssumption := entity.DependencyAssumptionFails + if probability >= 0.5 { + preferredAssumption = entity.DependencyAssumptionSucceeds + } + stream.base[i] = entity.PathDependency{ + Batch: dependencyID, + Assumption: preferredAssumption, + } + stream.optimalLogProbability += math.Log(preferredProbability) + + variable := dependencyVariable{dependencyIndex: i} + if nonPreferredProbability == 0 { + variable.zeroProbability = true + } else { + variable.flipCost = math.Log(preferredProbability / nonPreferredProbability) + } + stream.variables = append(stream.variables, variable) + } + + sort.SliceStable(stream.variables, func(i, j int) bool { + left, right := stream.variables[i], stream.variables[j] + if left.zeroProbability != right.zeroProbability { + return !left.zeroProbability + } + if left.flipCost != right.flipCost { + return left.flipCost < right.flipCost + } + return left.dependencyIndex < right.dependencyIndex + }) + + heap.Init(&stream.subsets) + if len(stream.variables) > 0 { + entry := flipEntry{ + lastIndex: 0, + flipped: []int{0}, + sequence: stream.takeSequence(), + } + entry.add(stream.variables[0]) + heap.Push(&stream.subsets, entry) + } + return stream +} + +func (s *pathStream) takeSequence() uint64 { + sequence := s.nextSequence + s.nextSequence++ + return sequence +} + +type rankedCandidate struct { + candidate entity.CandidatePath + logProbability float64 +} + +func (s *pathStream) next() (rankedCandidate, bool) { + if s.emitOptimal { + s.emitOptimal = false + return s.build(nil, flipCost{}), true + } + if s.subsets.Len() == 0 { + return rankedCandidate{}, false + } + + entry := heap.Pop(&s.subsets).(flipEntry) + j := entry.lastIndex + if j+1 < len(s.variables) { + extend := flipEntry{ + cost: entry.cost, + lastIndex: j + 1, + flipped: appendCopy(entry.flipped, j+1), + sequence: s.takeSequence(), + } + extend.add(s.variables[j+1]) + heap.Push(&s.subsets, extend) + + swap := flipEntry{ + cost: entry.cost, + lastIndex: j + 1, + flipped: replaceLastCopy(entry.flipped, j+1), + sequence: s.takeSequence(), + } + swap.remove(s.variables[j]) + swap.add(s.variables[j+1]) + heap.Push(&s.subsets, swap) + } + return s.build(entry.flipped, entry.cost), true +} + +func appendCopy(values []int, value int) []int { + result := make([]int, len(values)+1) + copy(result, values) + result[len(values)] = value + return result +} + +func replaceLastCopy(values []int, value int) []int { + result := make([]int, len(values)) + copy(result, values) + result[len(result)-1] = value + return result +} + +func (s *pathStream) build(flipped []int, cost flipCost) rankedCandidate { + dependencies := make([]entity.PathDependency, len(s.base)) + copy(dependencies, s.base) + for _, variableIndex := range flipped { + dependencyIndex := s.variables[variableIndex].dependencyIndex + dependencies[dependencyIndex].Assumption = opposite(dependencies[dependencyIndex].Assumption) + } + + logProbability := s.optimalLogProbability - cost.finite + rankingScore := math.Exp(logProbability) + if cost.zeroProbabilityFlips > 0 { + logProbability = math.Inf(-1) + rankingScore = 0 + } + return rankedCandidate{ + candidate: entity.CandidatePath{ + Path: entity.SpeculationPath{ + Head: s.head, + Dependencies: dependencies, + }, + RankingScore: rankingScore, + }, + logProbability: logProbability, + } +} + +func opposite(assumption entity.DependencyAssumption) entity.DependencyAssumption { + if assumption == entity.DependencyAssumptionSucceeds { + return entity.DependencyAssumptionFails + } + return entity.DependencyAssumptionSucceeds +} + +// flipCost separates finite log penalties from flips whose probability is +// exactly zero. Keeping a count avoids undefined Inf-Inf arithmetic in the +// extend/swap enumeration. +type flipCost struct { + finite float64 + zeroProbabilityFlips int +} + +type flipEntry struct { + cost flipCost + lastIndex int + flipped []int + sequence uint64 +} + +func (e *flipEntry) add(variable dependencyVariable) { + if variable.zeroProbability { + e.cost.zeroProbabilityFlips++ + return + } + e.cost.finite += variable.flipCost +} + +func (e *flipEntry) remove(variable dependencyVariable) { + if variable.zeroProbability { + e.cost.zeroProbabilityFlips-- + return + } + e.cost.finite -= variable.flipCost +} + +type flipHeap []flipEntry + +var _ heap.Interface = (*flipHeap)(nil) + +func (h flipHeap) Len() int { return len(h) } + +func (h flipHeap) Less(i, j int) bool { + left, right := h[i], h[j] + if left.cost.zeroProbabilityFlips != right.cost.zeroProbabilityFlips { + return left.cost.zeroProbabilityFlips < right.cost.zeroProbabilityFlips + } + if left.cost.finite != right.cost.finite { + return left.cost.finite < right.cost.finite + } + return left.sequence < right.sequence +} + +func (h flipHeap) Swap(i, j int) { + h[i], h[j] = h[j], h[i] +} + +func (h *flipHeap) Push(value any) { + *h = append(*h, value.(flipEntry)) +} + +func (h *flipHeap) Pop() any { + old := *h + last := len(old) - 1 + value := old[last] + *h = old[:last] + return value +} + +// candidateIterator performs a k-way merge of the per-head ordered streams. +type candidateIterator struct { + candidates candidateHeap +} + +var _ generator.Iterator = (*candidateIterator)(nil) + +func (i *candidateIterator) Next(ctx context.Context) (entity.CandidatePath, bool, error) { + if err := ctx.Err(); err != nil { + return entity.CandidatePath{}, false, err + } + if i.candidates.Len() == 0 { + return entity.CandidatePath{}, false, nil + } + + item := heap.Pop(&i.candidates).(candidateItem) + if next, ok := item.stream.next(); ok { + heap.Push(&i.candidates, candidateItem{ + rankedCandidate: next, + pathID: next.candidate.Path.ID(), + stream: item.stream, + }) + } + return item.candidate, true, nil +} + +type candidateItem struct { + rankedCandidate + pathID string + stream *pathStream +} + +// candidateHeap is a max-heap by internal log probability. RankingScore can +// underflow to zero for deep paths, so it is not used to preserve ordering. +type candidateHeap []candidateItem + +var _ heap.Interface = (*candidateHeap)(nil) + +func (h candidateHeap) Len() int { return len(h) } + +func (h candidateHeap) Less(i, j int) bool { + left, right := h[i], h[j] + if left.logProbability != right.logProbability { + return left.logProbability > right.logProbability + } + if left.candidate.Path.Head != right.candidate.Path.Head { + return left.candidate.Path.Head < right.candidate.Path.Head + } + return left.pathID < right.pathID +} + +func (h candidateHeap) Swap(i, j int) { + h[i], h[j] = h[j], h[i] +} + +func (h *candidateHeap) Push(value any) { + *h = append(*h, value.(candidateItem)) +} + +func (h *candidateHeap) Pop() any { + old := *h + last := len(old) - 1 + value := old[last] + *h = old[:last] + return value +} diff --git a/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go new file mode 100644 index 00000000..2177e873 --- /dev/null +++ b/submitqueue/extension/speculation/generator/bestfirst/bestfirst_test.go @@ -0,0 +1,571 @@ +// Copyright (c) 2026 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" + "sort" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/submitqueue/submitqueue/entity" +) + +type scorerFunc func(context.Context, entity.Batch) (float64, error) + +func (f scorerFunc) Score(ctx context.Context, batch entity.Batch) (float64, error) { + return f(ctx, batch) +} + +type recordingScorer struct { + scores map[string]float64 + calls map[string]int +} + +func (s *recordingScorer) Score(_ context.Context, batch entity.Batch) (float64, error) { + s.calls[batch.ID]++ + return s.scores[batch.ID], nil +} + +func drain(t *testing.T, it interface { + Next(context.Context) (entity.CandidatePath, bool, error) +}) []entity.CandidatePath { + t.Helper() + var candidates []entity.CandidatePath + for { + candidate, ok, err := it.Next(context.Background()) + require.NoError(t, err) + if !ok { + return candidates + } + candidates = append(candidates, candidate) + } +} + +func pathKey(path entity.SpeculationPath) string { + parts := make([]string, 0, len(path.Dependencies)+1) + parts = append(parts, path.Head) + for _, dependency := range path.Dependencies { + parts = append(parts, dependency.Batch+"="+string(dependency.Assumption)) + } + return strings.Join(parts, "|") +} + +func assumptionFor(t *testing.T, path entity.SpeculationPath, batchID string) entity.DependencyAssumption { + t.Helper() + for _, dependency := range path.Dependencies { + if dependency.Batch == batchID { + return dependency.Assumption + } + } + t.Fatalf("dependency %q not found in path %+v", batchID, path) + return entity.DependencyAssumptionUnknown +} + +func TestNew_RequiresScorer(t *testing.T) { + assert.Panics(t, func() { + New(nil) + }) +} + +func TestGenerate_ExampleOrdersCandidatesAcrossDisconnectedHeads(t *testing.T) { + s := &recordingScorer{ + scores: map[string]float64{ + "A": 0.8, + "B": 0.6, + "E": 0.3, + }, + calls: make(map[string]int), + } + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "B", State: entity.BatchStateCreated}, + {ID: "C", State: entity.BatchStateSpeculating, Dependencies: []string{"A", "B"}}, + {ID: "E", State: entity.BatchStateCreated}, + {ID: "F", State: entity.BatchStateSpeculating, Dependencies: []string{"E"}}, + } + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + got := drain(t, it) + + require.Len(t, got, 6) + assert.Equal(t, []string{ + "F|E=fails", + "C|A=succeeds|B=succeeds", + "C|A=succeeds|B=fails", + "F|E=succeeds", + "C|A=fails|B=succeeds", + "C|A=fails|B=fails", + }, []string{ + pathKey(got[0].Path), + pathKey(got[1].Path), + pathKey(got[2].Path), + pathKey(got[3].Path), + pathKey(got[4].Path), + pathKey(got[5].Path), + }) + assert.InDelta(t, 0.70, got[0].RankingScore, 1e-12) + assert.InDelta(t, 0.48, got[1].RankingScore, 1e-12) + assert.InDelta(t, 0.32, got[2].RankingScore, 1e-12) + assert.InDelta(t, 0.30, got[3].RankingScore, 1e-12) + assert.InDelta(t, 0.12, got[4].RankingScore, 1e-12) + assert.InDelta(t, 0.08, got[5].RankingScore, 1e-12) + assert.Equal(t, map[string]int{"A": 1, "B": 1, "E": 1}, s.calls) +} + +func TestGenerate_ResolvedDependenciesBecomeFixedAssumptions(t *testing.T) { + s := &recordingScorer{ + scores: map[string]float64{"U": 0.7}, + calls: make(map[string]int), + } + batches := []entity.Batch{ + {ID: "S", State: entity.BatchStateSucceeded}, + {ID: "F", State: entity.BatchStateFailed}, + {ID: "X", State: entity.BatchStateCancelled}, + {ID: "U", State: entity.BatchStateMerging}, + { + ID: "H", + State: entity.BatchStateSpeculating, + Dependencies: []string{"S", "F", "X", "U"}, + }, + } + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + got := drain(t, it) + + require.Len(t, got, 2) + for _, candidate := range got { + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(t, candidate.Path, "S")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(t, candidate.Path, "F")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(t, candidate.Path, "X")) + } + assert.Equal(t, entity.DependencyAssumptionSucceeds, assumptionFor(t, got[0].Path, "U")) + assert.Equal(t, entity.DependencyAssumptionFails, assumptionFor(t, got[1].Path, "U")) + assert.InDelta(t, 0.7, got[0].RankingScore, 1e-12) + assert.InDelta(t, 0.3, got[1].RankingScore, 1e-12) + assert.Equal(t, map[string]int{"U": 1}, s.calls) +} + +func TestGenerate_AllResolvedDependenciesYieldOneCertainPath(t *testing.T) { + s := scorerFunc(func(_ context.Context, batch entity.Batch) (float64, error) { + t.Fatalf("unexpected score call for %q", batch.ID) + return 0, nil + }) + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateSucceeded}, + {ID: "B", State: entity.BatchStateFailed}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"A", "B"}}, + } + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + got := drain(t, it) + + require.Len(t, got, 1) + assert.Equal(t, "H|A=succeeds|B=fails", pathKey(got[0].Path)) + assert.Equal(t, 1.0, got[0].RankingScore) +} + +func TestGenerate_OnlySpeculatingBatchesAreHeads(t *testing.T) { + s := scorerFunc(func(_ context.Context, _ entity.Batch) (float64, error) { + return 0.5, nil + }) + batches := []entity.Batch{ + {ID: "created", State: entity.BatchStateCreated}, + {ID: "merging", State: entity.BatchStateMerging}, + {ID: "succeeded", State: entity.BatchStateSucceeded}, + {ID: "head", State: entity.BatchStateSpeculating}, + } + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + got := drain(t, it) + + require.Len(t, got, 1) + assert.Equal(t, "head", got[0].Path.Head) + assert.Empty(t, got[0].Path.Dependencies) + assert.Equal(t, 1.0, got[0].RankingScore) +} + +func TestGenerate_SpeculatingBatchCanAlsoBeDependency(t *testing.T) { + s := &recordingScorer{ + scores: map[string]float64{"A": 0.8}, + calls: make(map[string]int), + } + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateSpeculating}, + {ID: "B", State: entity.BatchStateSpeculating, Dependencies: []string{"A"}}, + } + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + got := drain(t, it) + + require.Len(t, got, 3) + assert.Equal(t, "A", got[0].Path.Head) + assert.Equal(t, 1.0, got[0].RankingScore) + assert.Equal(t, "B|A=succeeds", pathKey(got[1].Path)) + assert.Equal(t, "B|A=fails", pathKey(got[2].Path)) + assert.Equal(t, map[string]int{"A": 1}, s.calls) +} + +func TestGenerate_ZeroAndOneProbabilitiesRemainExact(t *testing.T) { + s := &recordingScorer{ + scores: map[string]float64{"A": 1, "B": 0}, + calls: make(map[string]int), + } + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "B", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"A", "B"}}, + } + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + got := drain(t, it) + + require.Len(t, got, 4) + assert.Equal(t, "H|A=succeeds|B=fails", pathKey(got[0].Path)) + assert.Equal(t, 1.0, got[0].RankingScore) + seen := make(map[string]struct{}, len(got)) + for i, candidate := range got { + assert.False(t, math.IsNaN(candidate.RankingScore)) + if i > 0 { + assert.Equal(t, 0.0, candidate.RankingScore) + } + seen[pathKey(candidate.Path)] = struct{}{} + } + assert.Len(t, seen, 4) +} + +func TestGenerate_ExhaustiveResultsMatchBruteForce(t *testing.T) { + probabilities := map[string]float64{ + "A": 0.9, + "B": 0.7, + "C": 0.6, + "D": 0.2, + } + s := &recordingScorer{scores: probabilities, calls: make(map[string]int)} + dependencies := []string{"A", "B", "C", "D"} + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "B", State: entity.BatchStateCreated}, + {ID: "C", State: entity.BatchStateCreated}, + {ID: "D", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: dependencies}, + } + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + got := drain(t, it) + require.Len(t, got, 1< 0 { + assert.GreaterOrEqual(t, got[i-1].RankingScore, candidate.RankingScore) + } + } + assert.Len(t, seen, len(expected)) +} + +func TestGenerate_EqualProbabilitiesAreDeterministicAndExhaustive(t *testing.T) { + s := scorerFunc(func(_ context.Context, _ entity.Batch) (float64, error) { + return 0.5, nil + }) + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "B", State: entity.BatchStateCreated}, + {ID: "C", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"A", "B", "C"}}, + } + + generateKeys := func() []string { + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + candidates := drain(t, it) + keys := make([]string, len(candidates)) + for i, candidate := range candidates { + keys[i] = pathKey(candidate.Path) + assert.InDelta(t, 0.125, candidate.RankingScore, 1e-12) + } + return keys + } + + first := generateKeys() + second := generateKeys() + assert.Equal(t, first, second) + assert.Len(t, first, 8) + assert.Len(t, mapFromStrings(first), 8) +} + +func mapFromStrings(values []string) map[string]struct{} { + result := make(map[string]struct{}, len(values)) + for _, value := range values { + result[value] = struct{}{} + } + return result +} + +func TestGenerate_PreservesDependencyOrderAndInput(t *testing.T) { + s := scorerFunc(func(_ context.Context, _ entity.Batch) (float64, error) { + return 0.8, nil + }) + dependencies := []string{"B", "A"} + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "B", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: dependencies}, + } + original := append([]string(nil), dependencies...) + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + candidate, ok, err := it.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + + assert.Equal(t, original, dependencies) + require.Len(t, candidate.Path.Dependencies, 2) + assert.Equal(t, "B", candidate.Path.Dependencies[0].Batch) + assert.Equal(t, "A", candidate.Path.Dependencies[1].Batch) +} + +func TestGenerate_LargeDependencySetIsLazy(t *testing.T) { + const dependencyCount = 50 + scores := make(map[string]float64, dependencyCount) + batches := make([]entity.Batch, 0, dependencyCount+1) + dependencies := make([]string, dependencyCount) + for i := range dependencies { + id := fmt.Sprintf("D%02d", i) + dependencies[i] = id + scores[id] = 0.8 + batches = append(batches, entity.Batch{ID: id, State: entity.BatchStateCreated}) + } + batches = append(batches, entity.Batch{ + ID: "H", + State: entity.BatchStateSpeculating, + Dependencies: dependencies, + }) + s := &recordingScorer{scores: scores, calls: make(map[string]int)} + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + var previous = math.Inf(1) + for range 32 { + candidate, ok, err := it.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + assert.LessOrEqual(t, candidate.RankingScore, previous) + previous = candidate.RankingScore + } + assert.Len(t, s.calls, dependencyCount) +} + +func TestGenerate_EmptyIterator(t *testing.T) { + s := scorerFunc(func(_ context.Context, batch entity.Batch) (float64, error) { + t.Fatalf("unexpected score call for %q", batch.ID) + return 0, nil + }) + + for _, batches := range [][]entity.Batch{ + nil, + {{ID: "A", State: entity.BatchStateCreated}}, + } { + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + candidate, ok, err := it.Next(context.Background()) + require.NoError(t, err) + assert.False(t, ok) + assert.Equal(t, entity.CandidatePath{}, candidate) + } +} + +func TestGenerate_RejectsMalformedSnapshots(t *testing.T) { + s := scorerFunc(func(_ context.Context, _ entity.Batch) (float64, error) { + return 0.5, nil + }) + tests := []struct { + name string + batches []entity.Batch + }{ + { + name: "empty batch ID", + batches: []entity.Batch{{State: entity.BatchStateSpeculating}}, + }, + { + name: "duplicate batch ID", + batches: []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "A", State: entity.BatchStateSpeculating}, + }, + }, + { + name: "empty dependency ID", + batches: []entity.Batch{ + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{""}}, + }, + }, + { + name: "self dependency", + batches: []entity.Batch{ + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"H"}}, + }, + }, + { + name: "duplicate dependency", + batches: []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"A", "A"}}, + }, + }, + { + name: "missing dependency", + batches: []entity.Batch{ + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"missing"}}, + }, + }, + { + name: "unknown dependency state", + batches: []entity.Batch{ + {ID: "A", State: entity.BatchStateUnknown}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"A"}}, + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := New(s).Generate(context.Background(), tt.batches) + require.Error(t, err) + }) + } +} + +func TestGenerate_RejectsInvalidProbabilities(t *testing.T) { + tests := []struct { + name string + probability float64 + }{ + {name: "negative", probability: -0.1}, + {name: "above one", probability: 1.1}, + {name: "NaN", probability: math.NaN()}, + {name: "positive infinity", probability: math.Inf(1)}, + {name: "negative infinity", probability: math.Inf(-1)}, + } + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"A"}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := scorerFunc(func(_ context.Context, _ entity.Batch) (float64, error) { + return tt.probability, nil + }) + _, err := New(s).Generate(context.Background(), batches) + require.Error(t, err) + }) + } +} + +func TestGenerate_PropagatesScorerError(t *testing.T) { + sentinel := errors.New("score failed") + s := scorerFunc(func(_ context.Context, _ entity.Batch) (float64, error) { + return 0, sentinel + }) + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"A"}}, + } + + _, err := New(s).Generate(context.Background(), batches) + require.ErrorIs(t, err, sentinel) +} + +func TestGenerate_AndNextHonorContextCancellation(t *testing.T) { + s := scorerFunc(func(_ context.Context, _ entity.Batch) (float64, error) { + return 0.8, nil + }) + batches := []entity.Batch{ + {ID: "A", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"A"}}, + } + + cancelled, cancel := context.WithCancel(context.Background()) + cancel() + _, err := New(s).Generate(cancelled, batches) + require.ErrorIs(t, err, context.Canceled) + + it, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + _, ok, err := it.Next(cancelled) + require.ErrorIs(t, err, context.Canceled) + assert.False(t, ok) + + candidate, ok, err := it.Next(context.Background()) + require.NoError(t, err) + require.True(t, ok) + assert.Equal(t, "H|A=succeeds", pathKey(candidate.Path)) +} + +func TestGenerate_ScoresDependenciesInStableOrder(t *testing.T) { + var calls []string + s := scorerFunc(func(_ context.Context, batch entity.Batch) (float64, error) { + calls = append(calls, batch.ID) + return 0.5, nil + }) + batches := []entity.Batch{ + {ID: "C", State: entity.BatchStateCreated}, + {ID: "A", State: entity.BatchStateCreated}, + {ID: "B", State: entity.BatchStateCreated}, + {ID: "H", State: entity.BatchStateSpeculating, Dependencies: []string{"C", "A", "B"}}, + } + + _, err := New(s).Generate(context.Background(), batches) + require.NoError(t, err) + assert.True(t, sort.StringsAreSorted(calls)) + assert.Equal(t, []string{"A", "B", "C"}, calls) +} diff --git a/submitqueue/extension/speculation/generator/generator.go b/submitqueue/extension/speculation/generator/generator.go new file mode 100644 index 00000000..9a411051 --- /dev/null +++ b/submitqueue/extension/speculation/generator/generator.go @@ -0,0 +1,47 @@ +// Copyright (c) 2026 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-path stream used by the default +// speculation implementation. +package generator + +//go:generate mockgen -source=generator.go -destination=mock/generator_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// Generator opens a best-first stream of coherent speculation candidates over +// one queue snapshot. +type Generator interface { + // Generate creates an Iterator over candidates for batches in + // BatchStateSpeculating. The input must include every batch referenced by a + // candidate head's Dependencies field, including terminal dependencies whose + // outcomes constrain the generated assumptions. + // + // The returned Iterator owns its state and is not safe for concurrent use. + // An empty snapshot or a snapshot without Speculating heads returns an empty + // Iterator rather than an error. + Generate(ctx context.Context, batches []entity.Batch) (Iterator, error) +} + +// Iterator is a pull-based, probability-ordered stream of candidate paths. +type Iterator interface { + // Next returns the next candidate. ok is false when the stream is exhausted; + // exhaustion is an expected result, not an error. Implementations must leave + // the stream unconsumed when ctx is already cancelled. + Next(ctx context.Context) (candidate 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..22740a6b --- /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 +} + +// Generate mocks base method. +func (m *MockGenerator) Generate(ctx context.Context, batches []entity.Batch) (generator.Iterator, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Generate", ctx, batches) + ret0, _ := ret[0].(generator.Iterator) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Generate indicates an expected call of Generate. +func (mr *MockGeneratorMockRecorder) Generate(ctx, batches any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Generate", reflect.TypeOf((*MockGenerator)(nil).Generate), ctx, batches) +} + +// MockIterator is a mock of Iterator interface. +type MockIterator struct { + ctrl *gomock.Controller + recorder *MockIteratorMockRecorder + isgomock struct{} +} + +// MockIteratorMockRecorder is the mock recorder for MockIterator. +type MockIteratorMockRecorder struct { + mock *MockIterator +} + +// NewMockIterator creates a new mock instance. +func NewMockIterator(ctrl *gomock.Controller) *MockIterator { + mock := &MockIterator{ctrl: ctrl} + mock.recorder = &MockIteratorMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockIterator) EXPECT() *MockIteratorMockRecorder { + return m.recorder +} + +// Next mocks base method. +func (m *MockIterator) 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 *MockIteratorMockRecorder) Next(ctx any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Next", reflect.TypeOf((*MockIterator)(nil).Next), ctx) +}