Refactor "Update batch AI to use new Selector" (#3428)

Author: leszko

server/ai_session.go

@@ -191,14 +191,12 @@ func NewAISessionSelector(ctx context.Context, cap core.Capability, modelID stri
 	if cap == core.Capability_LiveVideoToVideo {
 		// For Realtime Video AI, we don't use any features of MinLSSelector (preferring known sessions, etc.),
 		// We always select a fresh session which has the lowest initial latency
-		useInitialLatencyToSort := true
-		warmSel = NewSelector(stakeRdr, node.SelectionAlgorithm, node.OrchPerfScore, warmCaps, useInitialLatencyToSort)
-		coldSel = NewSelector(stakeRdr, node.SelectionAlgorithm, node.OrchPerfScore, coldCaps, useInitialLatencyToSort)
+		warmSel = NewSelector(stakeRdr, node.SelectionAlgorithm, node.OrchPerfScore, warmCaps)
+		coldSel = NewSelector(stakeRdr, node.SelectionAlgorithm, node.OrchPerfScore, coldCaps)
 	} else {
-		//sort sessions based on current latency score
-		useInitialLatencyToSort := false
-		warmSel = NewSelector(stakeRdr, node.SelectionAlgorithm, node.OrchPerfScore, warmCaps, useInitialLatencyToSort)
-		coldSel = NewSelector(stakeRdr, node.SelectionAlgorithm, node.OrchPerfScore, coldCaps, useInitialLatencyToSort)
+		// sort sessions based on current latency score
+		warmSel = NewSelectorOrderByLatencyScore(stakeRdr, node.SelectionAlgorithm, node.OrchPerfScore, warmCaps)
+		coldSel = NewSelectorOrderByLatencyScore(stakeRdr, node.SelectionAlgorithm, node.OrchPerfScore, coldCaps)
 	}
 
 	warmPool := NewAISessionPool(warmSel, suspender, penalty)

server/selection.go

@@ -94,49 +94,60 @@ func (r *storeStakeReader) Stakes(addrs []ethcommon.Address) (map[ethcommon.Addr
 type Selector struct {
 	sessions []*BroadcastSession
 
-	stakeRdr                stakeReader
-	selectionAlgorithm      common.SelectionAlgorithm
-	perfScore               *common.PerfScore
-	capabilities            common.CapabilityComparator
-	useInitialLatencyToSort bool
+	stakeRdr           stakeReader
+	selectionAlgorithm common.SelectionAlgorithm
+	perfScore          *common.PerfScore
+	capabilities       common.CapabilityComparator
+	sortCompFunc       func(sess1, sess2 *BroadcastSession) bool
 }
 
-func NewSelector(stakeRdr stakeReader, selectionAlgorithm common.SelectionAlgorithm, perfScore *common.PerfScore, capabilities common.CapabilityComparator, useInitialLatencyToSort bool) *Selector {
+func NewSelector(stakeRdr stakeReader, selectionAlgorithm common.SelectionAlgorithm, perfScore *common.PerfScore, capabilities common.CapabilityComparator) *Selector {
+	// By default, sort by initial latency
+	sortCompFunc := func(sess1, sess2 *BroadcastSession) bool {
+		return sess1.InitialLatency < sess2.InitialLatency
+	}
 	return &Selector{
-		stakeRdr:                stakeRdr,
-		selectionAlgorithm:      selectionAlgorithm,
-		perfScore:               perfScore,
-		capabilities:            capabilities,
-		useInitialLatencyToSort: useInitialLatencyToSort,
+		stakeRdr:           stakeRdr,
+		selectionAlgorithm: selectionAlgorithm,
+		perfScore:          perfScore,
+		capabilities:       capabilities,
+		sortCompFunc:       sortCompFunc,
+	}
+}
+
+func NewSelectorOrderByLatencyScore(stakeRdr stakeReader, selectionAlgorithm common.SelectionAlgorithm, perfScore *common.PerfScore, capabilities common.CapabilityComparator) *Selector {
+	sortCompFunc := func(sess1, sess2 *BroadcastSession) bool {
+		return sess1.LatencyScore < sess2.LatencyScore
+	}
+	return &Selector{
+		stakeRdr:           stakeRdr,
+		selectionAlgorithm: selectionAlgorithm,
+		perfScore:          perfScore,
+		capabilities:       capabilities,
+		sortCompFunc:       sortCompFunc,
 	}
 }
 
 func (s *Selector) Add(sessions []*BroadcastSession) {
 	s.sessions = append(s.sessions, sessions...)
-	s.sortByLatency()
+	s.sort()
 }
 
 func (s *Selector) Complete(sess *BroadcastSession) {
 	s.sessions = append(s.sessions, sess)
-	s.sortByLatency()
-}
-
-func (s *Selector) sortByLatency() {
-	if s.useInitialLatencyToSort {
-		sort.Slice(s.sessions, func(i, j int) bool {
-			return s.sessions[i].InitialLatency < s.sessions[j].InitialLatency
-		})
-	} else {
-		sort.Slice(s.sessions, func(i, j int) bool {
-			return s.sessions[i].LatencyScore < s.sessions[j].LatencyScore
-		})
-	}
+	s.sort()
+}
+
+func (s *Selector) sort() {
+	sort.Slice(s.sessions, func(i, j int) bool {
+		return s.sortCompFunc(s.sessions[i], s.sessions[j])
+	})
 }
 
 func (s *Selector) Select(ctx context.Context) *BroadcastSession {
 	availableOrchestrators := toOrchestrators(s.sessions)
 	sess := s.selectUnknownSession(ctx)
-	s.sortByLatency()
+	s.sort()
 	clog.V(common.DEBUG).Infof(ctx, "Selected orchestrator %s from available list: %v", toOrchestrator(sess), availableOrchestrators)
 	return sess
 }

server/selection_test.go

@@ -162,7 +162,7 @@ func TestSelector_Select(t *testing.T) {
 	assert := assert.New(t)
 
 	// given
-	sel := NewSelector(nil, stubSelectionAlgorithm{}, nil, nil, true)
+	sel := NewSelector(nil, stubSelectionAlgorithm{}, nil, nil)
 	sessions := []*BroadcastSession{
 		{PMSessionID: "session-1", InitialLatency: 400 * time.Millisecond},
 		{PMSessionID: "session-2", InitialLatency: 200 * time.Millisecond},
@@ -185,7 +185,7 @@ func TestSelector_CompleteAndSelect(t *testing.T) {
 	assert := assert.New(t)
 
 	// given
-	sel := NewSelector(nil, stubSelectionAlgorithm{}, nil, nil, true)
+	sel := NewSelector(nil, stubSelectionAlgorithm{}, nil, nil)
 	sessions := []*BroadcastSession{
 		{PMSessionID: "session-1", InitialLatency: 400 * time.Millisecond},
 		{PMSessionID: "session-2", InitialLatency: 200 * time.Millisecond},
@@ -213,7 +213,7 @@ func TestSelector_Size(t *testing.T) {
 	assert := assert.New(t)
 
 	// given
-	sel := NewSelector(nil, stubSelectionAlgorithm{}, nil, nil, true)
+	sel := NewSelector(nil, stubSelectionAlgorithm{}, nil, nil)
 	sessions := []*BroadcastSession{
 		{PMSessionID: "session-1", InitialLatency: 400 * time.Millisecond},
 		{PMSessionID: "session-2", InitialLatency: 200 * time.Millisecond},
@@ -238,11 +238,10 @@ func TestSelector_Size(t *testing.T) {
 	assert.Nil(sel.Select(context.Background()))
 }
 
-func TestSelector_SortByLatency(t *testing.T) {
+func TestSelector_SortByInitialLatency(t *testing.T) {
 	assert := assert.New(t)
 
-	// sort by initial latency
-	sel := NewSelector(nil, stubSelectionAlgorithm{}, nil, nil, true)
+	sel := NewSelector(nil, stubSelectionAlgorithm{}, nil, nil)
 	sessions := []*BroadcastSession{
 		{PMSessionID: "session-1", InitialLatency: 400 * time.Millisecond},
 		{PMSessionID: "session-2", InitialLatency: 200 * time.Millisecond},
@@ -253,10 +252,13 @@ func TestSelector_SortByLatency(t *testing.T) {
 	assert.Equal("session-2", sel.sessions[0].PMSessionID)
 	assert.Equal("session-1", sel.sessions[1].PMSessionID)
 	assert.Equal("session-3", sel.sessions[2].PMSessionID)
+}
+
+func TestSelector_SortByLatencyScore(t *testing.T) {
+	assert := assert.New(t)
 
-	// sort by initial latency
-	sel = NewSelector(nil, stubSelectionAlgorithm{}, nil, nil, false)
-	sessions = []*BroadcastSession{
+	sel := NewSelectorOrderByLatencyScore(nil, stubSelectionAlgorithm{}, nil, nil)
+	sessions := []*BroadcastSession{
 		{PMSessionID: "session-1", InitialLatency: 400 * time.Millisecond, LatencyScore: 0.001},
 		{PMSessionID: "session-2", InitialLatency: 200 * time.Millisecond, LatencyScore: 0.01},
 		{PMSessionID: "session-3", InitialLatency: 600 * time.Millisecond, LatencyScore: 0.08},