Add node and queue success probability estimation (#3401)

* Add node and queue success probability estimation

* Docs, cleanup

* Docs

* Docs

* Docs

* Docs

* Initialise new nodes (queues) to a 50% success prob

* Docs

* Docs

* Reduce default queue timeout

config/scheduler/config.yaml

@@ -139,4 +139,11 @@ scheduling:
       resolution: "1Mi"
   gangIdAnnotation: armadaproject.io/gangId
   gangCardinalityAnnotation: armadaproject.io/gangCardinality
+  failureEstimatorConfig:
+    nodeSuccessProbabilityCordonThreshold: 0.1
+    queueSuccessProbabilityCordonThreshold: 0.05
+    nodeCordonTimeout: "10m"
+    queueCordonTimeout: "1m"
+    nodeEquilibriumFailureRate: 0.0167 # 1 per minute.
+    queueEquilibriumFailureRate: 0.0167 # 1 per minute.
 

internal/armada/configuration/types.go

@@ -219,6 +219,8 @@ type SchedulingConfig struct {
 	ExecutorUpdateFrequency time.Duration
 	// Enable new preemption strategy.
 	EnableNewPreemptionStrategy bool
+	// Controls node and queue success probability estimation.
+	FailureEstimatorConfig FailureEstimatorConfig
 }
 
 const (
@@ -315,6 +317,18 @@ type PreemptionConfig struct {
 	PriorityClassNameOverride *string
 }
 
+// FailureEstimatorConfig contains config controlling node and queue success probability estimation.
+// See the internal/scheduler/failureestimator package for details.
+type FailureEstimatorConfig struct {
+	Disabled                               bool
+	NodeSuccessProbabilityCordonThreshold  float64
+	QueueSuccessProbabilityCordonThreshold float64
+	NodeCordonTimeout                      time.Duration
+	QueueCordonTimeout                     time.Duration
+	NodeEquilibriumFailureRate             float64
+	QueueEquilibriumFailureRate            float64
+}
+
 // TODO: we can probably just typedef this to map[string]string
 type PostgresConfig struct {
 	Connection map[string]string

internal/scheduler/failureestimator/failureestimator.go

@@ -0,0 +1,307 @@
+package failureestimator
+
+import (
+	"fmt"
+	"math"
+	"sync"
+	"time"
+
+	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
+
+	"github.com/armadaproject/armada/internal/common/armadaerrors"
+)
+
+const (
+	namespace = "armada"
+	subsystem = "scheduler"
+
+	// Floating point tolerance. Also used when applying limits to avoid divide-by-zero issues.
+	eps = 1e-15
+	// Assumed success probability of "good" nodes (queues) used when calculating step size.
+	healthySuccessProbability = 0.95
+)
+
+// FailureEstimator is a system for answering the following question:
+// "What's the probability of a job from queue Q completing successfully when scheduled on node N?"
+// We assume the job may fail either because the job or node is faulty, and we assume these failures are independent.
+// Denote by
+// - P_q the probability of a job from queue q succeeding when running on a perfect node and
+// - P_n is the probability of a perfect job succeeding on node n.
+// The success probability of a job from queue q on node n is then Pr(p_q*p_n=1),
+// where p_q and p_n are drawn from Bernoulli distributions with parameter P_q and P_n, respectively.
+//
+// Now, the goal is to jointly estimate P_q and P_n for each queue and node using observed successes and failures.
+// The method used is statistical and only relies on knowing which queue a job belongs to and on which node it ran.
+// The intuition of the method is that:
+// - A job from a queue with many failures doesn't say much about the node; likely it's the job that's the problem.
+// - A job failing on a node with many failures doesn't say much about the job; likely it's the node that's the problem.
+// And vice versa.
+//
+// Specifically, we maximise the log-likelihood function of P_q and P_n using observed successes and failures.
+// This maximisation is performed using online gradient descent, where for each success or failure,
+// we update the corresponding P_q and P_n by taking a gradient step.
+// See New(...) for more details regarding step size.
+//
+// Finally, we exponentially decay P_q and P_N towards 1 over time,
+// such that nodes and queues for which we observe no failures appear to become healthier over time.
+// See New(...) function for details regarding decay.
+//
+// This module internally only maintains success probability estimates, as this makes the maths cleaner.
+// When exporting these via API calls we convert to failure probabilities as these are more intuitive to reason about.
+type FailureEstimator struct {
+	// Map from node (queue) name to the estimated success probability of that node (queue). For example:
+	// - successProbabilityByNode["myNode"] = 0.85]: estimated failure probability of a perfect job run on "myNode" is 15%.
+	// - successProbabilityByQueue["myQueue"] = 0.60]: estimated failure probability of a job from "myQueue" run on a perfect node is 40%.
+	successProbabilityByNode  map[string]float64
+	successProbabilityByQueue map[string]float64
+
+	// Success probability below which to consider nodes (jobs) unhealthy.
+	nodeSuccessProbabilityCordonThreshold  float64
+	queueSuccessProbabilityCordonThreshold float64
+
+	// Exponential decay factor controlling how quickly estimated node (queue) success probability decays towards 1.
+	// Computed from:
+	// - {node, queue}SuccessProbabilityCordonThreshold
+	// - {node, queue}CordonTimeout
+	nodeFailureProbabilityDecayRate  float64
+	queueFailureProbabilityDecayRate float64
+	timeOfLastDecay                  time.Time
+
+	// Gradient descent step size. Controls the extent to which new data affects successProbabilityBy{Node, Queue}.
+	// Computed from:
+	// - {node, queue}SuccessProbabilityCordonThreshold
+	// - {node, queue}FailureProbabilityDecayRate
+	// - {node, queue}EquilibriumFailureRate
+	nodeStepSize  float64
+	queueStepSize float64
+
+	// Prometheus metrics.
+	failureProbabilityByNodeDesc  *prometheus.Desc
+	failureProbabilityByQueueDesc *prometheus.Desc
+
+	// If true, this module is disabled.
+	disabled bool
+
+	// Mutex protecting the above fields.
+	// Prevents concurrent map modification issues when scraping metrics.
+	mu sync.Mutex
+}
+
+// New returns a new FailureEstimator. Parameters have the following meaning:
+// - {node, queue}SuccessProbabilityCordonThreshold: Success probability below which nodes (queues) are considered unhealthy.
+// - {node, queue}CordonTimeout: Amount of time for which nodes (queues) remain unhealthy in the absence of any job successes or failures for that node (queue).
+// - {node, queue}EquilibriumFailureRate: Job failures per second necessary for a node (queue) to remain unhealthy in the absence of any successes for that node (queue).
+func New(
+	nodeSuccessProbabilityCordonThreshold float64,
+	queueSuccessProbabilityCordonThreshold float64,
+	nodeCordonTimeout time.Duration,
+	queueCordonTimeout time.Duration,
+	nodeEquilibriumFailureRate float64,
+	queueEquilibriumFailureRate float64,
+) (*FailureEstimator, error) {
+	if nodeSuccessProbabilityCordonThreshold < 0 || nodeSuccessProbabilityCordonThreshold > 1 {
+		return nil, errors.WithStack(&armadaerrors.ErrInvalidArgument{
+			Name:    "nodeSuccessProbabilityCordonThreshold",
+			Value:   nodeSuccessProbabilityCordonThreshold,
+			Message: fmt.Sprintf("outside allowed range [0, 1]"),
+		})
+	}
+	if queueSuccessProbabilityCordonThreshold < 0 || queueSuccessProbabilityCordonThreshold > 1 {
+		return nil, errors.WithStack(&armadaerrors.ErrInvalidArgument{
+			Name:    "queueSuccessProbabilityCordonThreshold",
+			Value:   queueSuccessProbabilityCordonThreshold,
+			Message: fmt.Sprintf("outside allowed range [0, 1]"),
+		})
+	}
+	if nodeCordonTimeout < 0 {
+		return nil, errors.WithStack(&armadaerrors.ErrInvalidArgument{
+			Name:    "nodeCordonTimeout",
+			Value:   nodeCordonTimeout,
+			Message: fmt.Sprintf("outside allowed range [0, Inf)"),
+		})
+	}
+	if queueCordonTimeout < 0 {
+		return nil, errors.WithStack(&armadaerrors.ErrInvalidArgument{
+			Name:    "queueCordonTimeout",
+			Value:   queueCordonTimeout,
+			Message: fmt.Sprintf("outside allowed range [0, Inf)"),
+		})
+	}
+	if nodeEquilibriumFailureRate <= 0 {
+		return nil, errors.WithStack(&armadaerrors.ErrInvalidArgument{
+			Name:    "nodeEquilibriumFailureRate",
+			Value:   nodeEquilibriumFailureRate,
+			Message: fmt.Sprintf("outside allowed range (0, Inf)"),
+		})
+	}
+	if queueEquilibriumFailureRate <= 0 {
+		return nil, errors.WithStack(&armadaerrors.ErrInvalidArgument{
+			Name:    "queueEquilibriumFailureRate",
+			Value:   queueEquilibriumFailureRate,
+			Message: fmt.Sprintf("outside allowed range (0, Inf)"),
+		})
+	}
+
+	// Compute decay rates such that a node (queue) with success probability 0 will over {node, queue}CordonTimeout time
+	// decay to a success probability of {node, queue}SuccessProbabilityCordonThreshold.
+	nodeFailureProbabilityDecayRate := math.Exp(math.Log(1-nodeSuccessProbabilityCordonThreshold) / nodeCordonTimeout.Seconds())
+	queueFailureProbabilityDecayRate := math.Exp(math.Log(1-queueSuccessProbabilityCordonThreshold) / queueCordonTimeout.Seconds())
+
+	// Compute step size such that a node (queue) with success probability {node, queue}SuccessProbabilityCordonThreshold
+	// for which we observe 0 successes and {node, queue}EquilibriumFailureRate failures per second from "good" nodes (queues)
+	// will remain at exactly {node, queue}SuccessProbabilityCordonThreshold success probability.
+	dNodeSuccessProbability := healthySuccessProbability / (1 - nodeSuccessProbabilityCordonThreshold*healthySuccessProbability)
+	dQueueSuccessProbability := healthySuccessProbability / (1 - queueSuccessProbabilityCordonThreshold*healthySuccessProbability)
+	nodeStepSize := (1 - nodeSuccessProbabilityCordonThreshold - (1-nodeSuccessProbabilityCordonThreshold)*nodeFailureProbabilityDecayRate) / dNodeSuccessProbability / nodeEquilibriumFailureRate
+	queueStepSize := (1 - queueSuccessProbabilityCordonThreshold - (1-queueSuccessProbabilityCordonThreshold)*queueFailureProbabilityDecayRate) / dQueueSuccessProbability / queueEquilibriumFailureRate
+
+	return &FailureEstimator{
+		successProbabilityByNode:               make(map[string]float64, 1024),
+		successProbabilityByQueue:              make(map[string]float64, 128),
+		nodeSuccessProbabilityCordonThreshold:  nodeSuccessProbabilityCordonThreshold,
+		queueSuccessProbabilityCordonThreshold: queueSuccessProbabilityCordonThreshold,
+		nodeFailureProbabilityDecayRate:        nodeFailureProbabilityDecayRate,
+		queueFailureProbabilityDecayRate:       queueFailureProbabilityDecayRate,
+		timeOfLastDecay:                        time.Now(),
+		nodeStepSize:                           nodeStepSize,
+		queueStepSize:                          queueStepSize,
+		failureProbabilityByNodeDesc: prometheus.NewDesc(
+			fmt.Sprintf("%s_%s_node_failure_probability", namespace, subsystem),
+			"Estimated per-node failure probability.",
+			[]string{"node"},
+			nil,
+		),
+		failureProbabilityByQueueDesc: prometheus.NewDesc(
+			fmt.Sprintf("%s_%s_queue_failure_probability", namespace, subsystem),
+			"Estimated per-queue failure probability.",
+			[]string{"queue"},
+			nil,
+		),
+	}, nil
+}
+
+func (fe *FailureEstimator) Disable(v bool) {
+	if fe == nil {
+		return
+	}
+	fe.disabled = v
+}
+
+func (fe *FailureEstimator) IsDisabled() bool {
+	if fe == nil {
+		return true
+	}
+	return fe.disabled
+}
+
+// Decay moves the success probabilities of nodes (queues) closer to 1, depending on the configured cordon timeout.
+// Periodically calling Decay() ensures nodes (queues) considered unhealthy are eventually considered healthy again,
+// even if we observe no successes for those nodes (queues).
+func (fe *FailureEstimator) Decay() {
+	fe.mu.Lock()
+	defer fe.mu.Unlock()
+	t := time.Now()
+	fe.decay(t.Sub(fe.timeOfLastDecay).Seconds())
+	fe.timeOfLastDecay = t
+	return
+}
+
+func (fe *FailureEstimator) decay(secondsSinceLastDecay float64) {
+	nodeFailureProbabilityDecay := math.Pow(fe.nodeFailureProbabilityDecayRate, secondsSinceLastDecay)
+	for k, v := range fe.successProbabilityByNode {
+		failureProbability := 1 - v
+		failureProbability *= nodeFailureProbabilityDecay
+		successProbability := 1 - failureProbability
+		fe.successProbabilityByNode[k] = applyBounds(successProbability)
+	}
+
+	queueFailureProbabilityDecay := math.Pow(fe.queueFailureProbabilityDecayRate, secondsSinceLastDecay)
+	for k, v := range fe.successProbabilityByQueue {
+		failureProbability := 1 - v
+		failureProbability *= queueFailureProbabilityDecay
+		successProbability := 1 - failureProbability
+		fe.successProbabilityByQueue[k] = applyBounds(successProbability)
+	}
+	return
+}
+
+// Update with success=false decreases the estimated success probability of the provided node and queue.
+// Calling with success=true increases the estimated success probability of the provided node and queue.
+// This update is performed by taking one gradient descent step.
+func (fe *FailureEstimator) Update(node, queue string, success bool) {
+	fe.mu.Lock()
+	defer fe.mu.Unlock()
+
+	// Assume that nodes (queues) we haven't seen before have a 50% success probability.
+	// Avoiding extreme values for new nodes (queues) helps avoid drastic changes to existing estimates.
+	nodeSuccessProbability, ok := fe.successProbabilityByNode[node]
+	if !ok {
+		nodeSuccessProbability = 0.5
+	}
+	queueSuccessProbability, ok := fe.successProbabilityByQueue[queue]
+	if !ok {
+		queueSuccessProbability = 0.5
+	}
+
+	dNodeSuccessProbability, dQueueSuccessProbability := fe.negLogLikelihoodGradient(nodeSuccessProbability, queueSuccessProbability, success)
+	nodeSuccessProbability -= fe.nodeStepSize * dNodeSuccessProbability
+	queueSuccessProbability -= fe.queueStepSize * dQueueSuccessProbability
+
+	fe.successProbabilityByNode[node] = applyBounds(nodeSuccessProbability)
+	fe.successProbabilityByQueue[queue] = applyBounds(queueSuccessProbability)
+}
+
+// applyBounds ensures values stay in the range [eps, 1-eps].
+// This to avoid divide-by-zero issues.
+func applyBounds(v float64) float64 {
+	if v < eps {
+		return eps
+	} else if v > 1.0-eps {
+		return 1.0 - eps
+	} else {
+		return v
+	}
+}
+
+// negLogLikelihoodGradient returns the gradient of the negated log-likelihood function with respect to P_q and P_n.
+func (fe *FailureEstimator) negLogLikelihoodGradient(nodeSuccessProbability, queueSuccessProbability float64, success bool) (float64, float64) {
+	if success {
+		dNodeSuccessProbability := -1 / nodeSuccessProbability
+		dQueueSuccessProbability := -1 / queueSuccessProbability
+		return dNodeSuccessProbability, dQueueSuccessProbability
+	} else {
+		dNodeSuccessProbability := queueSuccessProbability / (1 - nodeSuccessProbability*queueSuccessProbability)
+		dQueueSuccessProbability := nodeSuccessProbability / (1 - nodeSuccessProbability*queueSuccessProbability)
+		return dNodeSuccessProbability, dQueueSuccessProbability
+	}
+}
+
+func (fe *FailureEstimator) Describe(ch chan<- *prometheus.Desc) {
+	if fe.IsDisabled() {
+		return
+	}
+	ch <- fe.failureProbabilityByNodeDesc
+	ch <- fe.failureProbabilityByQueueDesc
+}
+
+func (fe *FailureEstimator) Collect(ch chan<- prometheus.Metric) {
+	if fe.IsDisabled() {
+		return
+	}
+	fe.mu.Lock()
+	defer fe.mu.Unlock()
+
+	// Report failure probability rounded to nearest multiple of 0.01.
+	// (As it's unlikely the estimate is accurate to within less than this.)
+	for k, v := range fe.successProbabilityByNode {
+		failureProbability := 1 - v
+		failureProbability = math.Round(failureProbability*100) / 100
+		ch <- prometheus.MustNewConstMetric(fe.failureProbabilityByNodeDesc, prometheus.GaugeValue, failureProbability, k)
+	}
+	for k, v := range fe.successProbabilityByQueue {
+		failureProbability := 1 - v
+		failureProbability = math.Round(failureProbability*100) / 100
+		ch <- prometheus.MustNewConstMetric(fe.failureProbabilityByQueueDesc, prometheus.GaugeValue, failureProbability, k)
+	}
+}