diff --git a/submitqueue/extension/speculation/pathscorer/probability/BUILD.bazel b/submitqueue/extension/speculation/pathscorer/probability/BUILD.bazel new file mode 100644 index 00000000..5fd08a1b --- /dev/null +++ b/submitqueue/extension/speculation/pathscorer/probability/BUILD.bazel @@ -0,0 +1,26 @@ +load("@rules_go//go:def.bzl", "go_library", "go_test") + +go_library( + name = "go_default_library", + srcs = ["probability.go"], + importpath = "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer/probability", + visibility = ["//visibility:public"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/speculation/pathscorer:go_default_library", + "//submitqueue/extension/storage:go_default_library", + ], +) + +go_test( + name = "go_default_test", + srcs = ["probability_test.go"], + embed = [":go_default_library"], + deps = [ + "//submitqueue/entity:go_default_library", + "//submitqueue/extension/storage/mock:go_default_library", + "@com_github_stretchr_testify//assert:go_default_library", + "@com_github_stretchr_testify//require:go_default_library", + "@org_uber_go_mock//gomock:go_default_library", + ], +) diff --git a/submitqueue/extension/speculation/pathscorer/probability/README.md b/submitqueue/extension/speculation/pathscorer/probability/README.md new file mode 100644 index 00000000..101fd6fb --- /dev/null +++ b/submitqueue/extension/speculation/pathscorer/probability/README.md @@ -0,0 +1,31 @@ +# Probability Path Scorer + +The probability `pathscorer.Scorer` scores every path in a batch's speculation tree as the probability that exactly that path's assumption holds: its base dependencies all land, and every other active dependency of the head fails to land. Concretely, a path's score is the head batch's probability, multiplied by the probability of each base dependency, multiplied by one minus the probability of each of the head's dependencies not in the base. + +Each batch's probability is resolved from its state, so a resolved outcome overrides the prediction. A batch that has succeeded contributes certainty 1; a batch that has failed or been cancelled contributes 0; a batch still in flight contributes `entity.Batch.Score`, the per-batch success probability set by the score stage, degrading to a neutral 0.5 when unscored so it neither helps nor hurts. A batch in the best-effort `cancelling` state keeps its prediction — it may still succeed. This state resolution is what redistributes score when the world changes: a dead dependency zeroes every path that built on it and boosts every path that excluded it by the full `(1 − p) = 1` factor, with no cross-path coupling needed. + +## Example + +Batch `C` has active dependencies `A` and `B`, with predicted success probabilities `p(C) = 0.8`, `p(A) = 0.9`, `p(B) = 0.6`, and a three-path tree: + +| Path | Base | Formula | Score | +|---|---|---|---| +| chain | `[A, B]` | 0.8 × 0.9 × 0.6 | **0.432** | +| drop-B | `[A]` | 0.8 × 0.9 × (1 − 0.6) | **0.288** | +| alone | `[]` | 0.8 × (1 − 0.9) × (1 − 0.6) | **0.032** | + +The chain path leads while everything looks healthy. Now `B`'s build fails (`p(B)` resolves to 0): + +| Path | Base | Formula | Score | +|---|---|---|---| +| chain | `[A, B]` | 0.8 × 0.9 × 0 | **0** | +| drop-B | `[A]` | 0.8 × 0.9 × 1 | **0.72** | +| alone | `[]` | 0.8 × 0.1 × 1 | **0.08** | + +The chain path — a bet on `B` landing — dies, and the drop-B path inherits its score mass. If `A` then lands (`p(A)` resolves to 1), drop-B rises to 0.8 and alone falls to 0. The selector and prioritizer rank on these scores, so builds follow the paths still consistent with reality. + +## Scope + +The model is deliberately simple: it treats dependency outcomes as independent, and does not weigh how long a batch has waited, its historical pass rate, or any correlation between siblings. It scores every path in the tree regardless of status, returning one path-ID-keyed score per path — the controller merges them into the tree and overwrites `Score` wholesale on every respeculate; nothing else about a path passes through the scorer. Folding resolved outcomes into path *status* (dead-pathing, cancellation) is the controller's reconcile job, not the scorer's; the scorer only makes the scores reflect them. + +Batches are read one at a time through the injected `storage.BatchStore`, matching the store's key/value contract, and each batch referenced by the tree is loaded at most once per `Score` call. A store error, including a missing batch, is returned wrapped and unclassified — scorer implementations, like all extensions, leave user/infra classification to the controller's error classifier. diff --git a/submitqueue/extension/speculation/pathscorer/probability/probability.go b/submitqueue/extension/speculation/pathscorer/probability/probability.go new file mode 100644 index 00000000..caf3b446 --- /dev/null +++ b/submitqueue/extension/speculation/pathscorer/probability/probability.go @@ -0,0 +1,142 @@ +// 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 probability provides a pathscorer.Scorer that scores each path as +// the probability that exactly that path's assumption holds: every base +// dependency lands and every non-base dependency of the head does not. Each +// dependency's probability is resolved from its batch's state — certainty +// (1 or 0) once the batch is terminal, its predicted build-success score +// while it is in flight — so a resolved dependency automatically kills the +// paths that bet against the outcome and boosts the paths consistent with +// it. Dependency outcomes are treated as independent; there is no adjustment +// for wait time, historical pass rate, or correlation between outcomes. +package probability + +import ( + "context" + "fmt" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/speculation/pathscorer" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +// neutralScore is the probability assigned to a non-terminal batch that has +// not yet been scored (entity.Batch.Score's zero value), so an unscored +// dependency neither helps nor hurts a path's score. +const neutralScore = 0.5 + +// probabilityScorer computes each path's Score as the probability that +// exactly its assumption holds: every base dependency lands and every other +// active dependency of the head fails to land — so this exact path, and no +// sibling, is the one that materializes. Batch probabilities come from +// batchProbability, which prefers a terminal outcome over the prediction. +// Folding resolved outcomes into path Status (dead-pathing, cancellation) +// remains the controller's reconcile concern; this scorer only recomputes +// Score. +type probabilityScorer struct { + // batches resolves a batch ID to its current state and predicted-success + // probability (entity.Batch.State / entity.Batch.Score) and, for the + // tree's head batch, its full active-dependency set + // (entity.Batch.Dependencies). + batches storage.BatchStore +} + +// New returns a pathscorer.Scorer implementing the path-probability +// heuristic, reading batch states and success probabilities from batches. +func New(batches storage.BatchStore) pathscorer.Scorer { + return &probabilityScorer{batches: batches} +} + +// batchProbability resolves the probability that b lands. A terminal state +// is certainty — 1 for succeeded, 0 for failed or cancelled — and overrides +// any prediction. Otherwise it is b.Score, the predicted build-success +// probability, degraded to a neutral 0.5 when the batch is unscored. +// Cancelling is deliberately not treated as terminal: cancellation is +// best-effort and the batch may still succeed, so the prediction stands +// until the state settles. +func batchProbability(b entity.Batch) float64 { + switch b.State { + case entity.BatchStateSucceeded: + return 1 + case entity.BatchStateFailed, entity.BatchStateCancelled: + return 0 + } + if b.Score == 0 { + return neutralScore + } + return b.Score +} + +// Score returns one PathScore per path in tree, regardless of Status, each +// naming its path by ID. A path's score is: +// +// p(head) * Π p(d) for d in path.Base * Π (1 - p(d)) for d in head.Dependencies but not in path.Base +// +// where p(b) is 1 when b has succeeded, 0 when b has failed or been +// cancelled, and otherwise b.Score (a neutral 0.5 when unscored) — see +// batchProbability. Each batch referenced by the tree is loaded from the +// store at most once. Any store error (including not-found) is returned +// wrapped, unclassified — extensions never classify errors as user or infra. +func (s *probabilityScorer) Score(ctx context.Context, tree entity.SpeculationTree) ([]entity.PathScore, error) { + if len(tree.Paths) == 0 { + return nil, nil + } + + head, err := s.batches.Get(ctx, tree.BatchID) + if err != nil { + return nil, fmt.Errorf("probability: get head batch %q: %w", tree.BatchID, err) + } + + probabilities := map[string]float64{} + probabilityOf := func(batchID string) (float64, error) { + if p, ok := probabilities[batchID]; ok { + return p, nil + } + b, err := s.batches.Get(ctx, batchID) + if err != nil { + return 0, fmt.Errorf("probability: get dependency batch %q: %w", batchID, err) + } + p := batchProbability(b) + probabilities[batchID] = p + return p, nil + } + + scores := make([]entity.PathScore, len(tree.Paths)) + for i, path := range tree.Paths { + inBase := make(map[string]bool, len(path.Path.Base)) + score := batchProbability(head) + for _, dep := range path.Path.Base { + inBase[dep] = true + p, err := probabilityOf(dep) + if err != nil { + return nil, err + } + score *= p + } + for _, dep := range head.Dependencies { + if inBase[dep] { + continue + } + p, err := probabilityOf(dep) + if err != nil { + return nil, err + } + score *= 1 - p + } + scores[i] = entity.PathScore{PathID: path.ID, Score: float32(score)} + } + + return scores, nil +} diff --git a/submitqueue/extension/speculation/pathscorer/probability/probability_test.go b/submitqueue/extension/speculation/pathscorer/probability/probability_test.go new file mode 100644 index 00000000..deb4d424 --- /dev/null +++ b/submitqueue/extension/speculation/pathscorer/probability/probability_test.go @@ -0,0 +1,257 @@ +// 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 probability + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/uber/submitqueue/submitqueue/entity" + storagemock "github.com/uber/submitqueue/submitqueue/extension/storage/mock" + "go.uber.org/mock/gomock" +) + +const headID = "q/batch/5" + +func TestScore_ProbabilityFormula(t *testing.T) { + tests := []struct { + name string + head entity.Batch + deps map[string]entity.Batch + path entity.SpeculationPath + want float32 + }{ + { + name: "healthy dep in base multiplies straight through", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1", Score: 0.9}}, + path: entity.SpeculationPath{Base: []string{"q/batch/1"}, Head: headID}, + want: 0.72, + }, + { + name: "shaky dep not in base contributes its failure probability", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1", Score: 0.9}}, + path: entity.SpeculationPath{Head: headID}, + want: 0.08, // 0.8 * (1 - 0.9) + }, + { + name: "unscored head degrades to neutral 0.5", + head: entity.Batch{ID: headID}, + deps: map[string]entity.Batch{}, + path: entity.SpeculationPath{Head: headID}, + want: 0.5, + }, + { + name: "unscored dep in base degrades to neutral 0.5", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1"}}, + path: entity.SpeculationPath{Base: []string{"q/batch/1"}, Head: headID}, + want: 0.4, // 0.8 * 0.5 + }, + { + name: "succeeded dep in base is certainty, overriding its prediction", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1", Score: 0.3, State: entity.BatchStateSucceeded}}, + path: entity.SpeculationPath{Base: []string{"q/batch/1"}, Head: headID}, + want: 0.8, // 0.8 * 1 + }, + { + name: "succeeded dep not in base zeroes the path betting against it", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1", Score: 0.3, State: entity.BatchStateSucceeded}}, + path: entity.SpeculationPath{Head: headID}, + want: 0, // 0.8 * (1 - 1) + }, + { + name: "failed dep in base zeroes the path built on it", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1", Score: 0.9, State: entity.BatchStateFailed}}, + path: entity.SpeculationPath{Base: []string{"q/batch/1"}, Head: headID}, + want: 0, // 0.8 * 0 + }, + { + name: "failed dep not in base boosts the path that excluded it", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1", Score: 0.9, State: entity.BatchStateFailed}}, + path: entity.SpeculationPath{Head: headID}, + want: 0.8, // 0.8 * (1 - 0) + }, + { + name: "cancelled dep in base zeroes the path built on it", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1", Score: 0.9, State: entity.BatchStateCancelled}}, + path: entity.SpeculationPath{Base: []string{"q/batch/1"}, Head: headID}, + want: 0, // 0.8 * 0 + }, + { + name: "cancelling dep keeps its prediction — cancellation is best-effort", + head: entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}}, + deps: map[string]entity.Batch{"q/batch/1": {ID: "q/batch/1", Score: 0.9, State: entity.BatchStateCancelling}}, + path: entity.SpeculationPath{Base: []string{"q/batch/1"}, Head: headID}, + want: 0.72, // 0.8 * 0.9 + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockBatchStore(ctrl) + store.EXPECT().Get(gomock.Any(), headID).Return(tt.head, nil) + for id, b := range tt.deps { + store.EXPECT().Get(gomock.Any(), id).Return(b, nil) + } + + tree := entity.SpeculationTree{ + BatchID: headID, + Paths: []entity.SpeculationPathInfo{{ID: "p/0", Path: tt.path, Status: entity.SpeculationPathStatusCandidate}}, + } + got, err := New(store).Score(context.Background(), tree) + require.NoError(t, err) + require.Len(t, got, 1) + assert.Equal(t, "p/0", got[0].PathID) + assert.InDelta(t, tt.want, got[0].Score, 0.0001) + }) + } +} + +// TestScore_DependencyFailureShiftsMassToSurvivingPaths pins the respeculate +// story end-to-end: with three paths over deps A and B, a failure of B zeroes +// the chain path that built on it and boosts the drop-B path by the full +// (1 - p(B)) = 1 factor — the score mass shifts to the paths consistent with +// the resolved outcome, with no cross-path coupling in the scorer. +func TestScore_DependencyFailureShiftsMassToSurvivingPaths(t *testing.T) { + depA := entity.Batch{ID: "q/batch/1", Score: 0.9} + depB := entity.Batch{ID: "q/batch/2", Score: 0.6} + head := entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{depA.ID, depB.ID}} + + tree := entity.SpeculationTree{ + BatchID: headID, + Paths: []entity.SpeculationPathInfo{ + {ID: "p/chain", Path: entity.SpeculationPath{Base: []string{depA.ID, depB.ID}, Head: headID}}, // full chain + {ID: "p/drop-b", Path: entity.SpeculationPath{Base: []string{depA.ID}, Head: headID}}, // drop B + {ID: "p/alone", Path: entity.SpeculationPath{Head: headID}}, // alone + }, + } + + score := func(t *testing.T, deps ...entity.Batch) []float32 { + ctrl := gomock.NewController(t) + store := storagemock.NewMockBatchStore(ctrl) + store.EXPECT().Get(gomock.Any(), headID).Return(head, nil) + for _, d := range deps { + store.EXPECT().Get(gomock.Any(), d.ID).Return(d, nil) + } + got, err := New(store).Score(context.Background(), tree) + require.NoError(t, err) + require.Len(t, got, 3) + scores := make([]float32, 3) + for i, ps := range got { + assert.Equal(t, tree.Paths[i].ID, ps.PathID) + scores[i] = ps.Score + } + return scores + } + + t.Run("before: both deps in flight, chain path leads", func(t *testing.T) { + scores := score(t, depA, depB) + assert.InDelta(t, 0.8*0.9*0.6, scores[0], 0.0001) // chain + assert.InDelta(t, 0.8*0.9*(1-0.6), scores[1], 0.0001) // drop B + assert.InDelta(t, 0.8*(1-0.9)*(1-0.6), scores[2], 0.0001) // alone + }) + + t.Run("after: B failed, drop-B path is boosted and chain dies", func(t *testing.T) { + failedB := depB + failedB.State = entity.BatchStateFailed + scores := score(t, depA, failedB) + assert.InDelta(t, 0, scores[0], 0.0001) // chain: bets on B landing + assert.InDelta(t, 0.8*0.9*1, scores[1], 0.0001) // drop B: (1 - 0) boost + assert.InDelta(t, 0.8*(1-0.9)*1, scores[2], 0.0001) // alone + }) +} + +func TestScore_MultiplePathsScoredIndependentlyAndDepsCachedOnce(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockBatchStore(ctrl) + + head := entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1", "q/batch/2"}} + dep1 := entity.Batch{ID: "q/batch/1", Score: 0.9} + dep2 := entity.Batch{ID: "q/batch/2", Score: 0.6} + + // Each batch is loaded at most once per Score call, even though both + // paths below reference dep1 and dep2. + store.EXPECT().Get(gomock.Any(), headID).Return(head, nil).Times(1) + store.EXPECT().Get(gomock.Any(), "q/batch/1").Return(dep1, nil).Times(1) + store.EXPECT().Get(gomock.Any(), "q/batch/2").Return(dep2, nil).Times(1) + + tree := entity.SpeculationTree{ + BatchID: headID, + Paths: []entity.SpeculationPathInfo{ + {ID: "p/0", Path: entity.SpeculationPath{Base: []string{"q/batch/1"}, Head: headID}}, + {ID: "p/1", Path: entity.SpeculationPath{Base: []string{"q/batch/2"}, Head: headID}}, + }, + } + + got, err := New(store).Score(context.Background(), tree) + require.NoError(t, err) + require.Len(t, got, 2) + assert.Equal(t, "p/0", got[0].PathID) + assert.InDelta(t, 0.8*0.9*(1-0.6), got[0].Score, 0.0001) + assert.Equal(t, "p/1", got[1].PathID) + assert.InDelta(t, 0.8*0.6*(1-0.9), got[1].Score, 0.0001) +} + +func TestScore_HeadStoreErrorPropagates(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockBatchStore(ctrl) + sentinel := errors.New("boom") + store.EXPECT().Get(gomock.Any(), headID).Return(entity.Batch{}, sentinel) + + tree := entity.SpeculationTree{ + BatchID: headID, + Paths: []entity.SpeculationPathInfo{{Path: entity.SpeculationPath{Head: headID}}}, + } + _, err := New(store).Score(context.Background(), tree) + require.ErrorIs(t, err, sentinel) +} + +func TestScore_DependencyStoreErrorPropagates(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockBatchStore(ctrl) + sentinel := errors.New("boom") + head := entity.Batch{ID: headID, Score: 0.8, Dependencies: []string{"q/batch/1"}} + store.EXPECT().Get(gomock.Any(), headID).Return(head, nil) + store.EXPECT().Get(gomock.Any(), "q/batch/1").Return(entity.Batch{}, sentinel) + + tree := entity.SpeculationTree{ + BatchID: headID, + Paths: []entity.SpeculationPathInfo{{Path: entity.SpeculationPath{Head: headID}}}, + } + _, err := New(store).Score(context.Background(), tree) + require.ErrorIs(t, err, sentinel) +} + +func TestScore_EmptyTreeSkipsStore(t *testing.T) { + ctrl := gomock.NewController(t) + store := storagemock.NewMockBatchStore(ctrl) + // No EXPECT() calls set up: the store must not be touched for an empty tree. + + tree := entity.SpeculationTree{BatchID: headID} + got, err := New(store).Score(context.Background(), tree) + require.NoError(t, err) + assert.Empty(t, got) +}