temp hold sessions with retryable errors

Author: ad-astra-video

server/ai_mediaserver.go

@@ -111,6 +111,7 @@ func aiMediaServerHandle[I, O any](ls *LivepeerServer, decoderFunc func(*I, *htt
 			node:        ls.LivepeerNode,
 			os:          drivers.NodeStorage.NewSession(requestID),
 			sessManager: ls.AISessionManager,
+			requestID:   requestID,
 		}
 
 		var req I
@@ -172,6 +173,7 @@ func (ls *LivepeerServer) ImageToVideo() http.Handler {
 			node:        ls.LivepeerNode,
 			os:          drivers.NodeStorage.NewSession(requestID),
 			sessManager: ls.AISessionManager,
+			requestID:   requestID,
 		}
 
 		if !async {
@@ -280,6 +282,7 @@ func (ls *LivepeerServer) LLM() http.Handler {
 			node:        ls.LivepeerNode,
 			os:          drivers.NodeStorage.NewSession(requestID),
 			sessManager: ls.AISessionManager,
+			requestID:   requestID,
 		}
 
 		start := time.Now()
@@ -554,6 +557,7 @@ func (ls *LivepeerServer) StartLiveVideo() http.Handler {
 			node:        ls.LivepeerNode,
 			os:          drivers.NodeStorage.NewSession(requestID),
 			sessManager: ls.AISessionManager,
+			requestID:   requestID,
 
 			liveParams: liveRequestParams{
 				segmentReader:          ssr,

server/ai_process.go

@@ -83,6 +83,7 @@ type aiRequestParams struct {
 	node        *core.LivepeerNode
 	os          drivers.OSSession
 	sessManager *AISessionManager
+	requestID   string
 
 	liveParams liveRequestParams
 }
@@ -1520,6 +1521,12 @@ func processAIRequest(ctx context.Context, params aiRequestParams, req interface
 			break
 		}
 
+		// Don't suspend the session if the error is a transient error.
+		if isRetryableError(err) {
+			params.sessManager.PutSessionOnHold(ctx, params.requestID, sess)
+			continue
+		}
+
 		// Suspend the session on other errors.
 		clog.Infof(ctx, "Error submitting request modelID=%v try=%v orch=%v err=%v", modelID, tries, sess.Transcoder(), err)
 		params.sessManager.Remove(ctx, sess) //TODO: Improve session selection logic for live-video-to-video
@@ -1543,6 +1550,25 @@ func processAIRequest(ctx context.Context, params aiRequestParams, req interface
 	return resp, nil
 }
 
+// isRetryableError checks if the error is a transient error that can be retried.
+var isRetryableError = retryableError
+
+func retryableError(err error) bool {
+	transientErrorMessages := []string{
+		"insufficient capacity",      // Caused by limitation in our current implementation.
+		"invalid ticket sendernonce", // Caused by gateway nonce mismatch.
+		"ticketparams expired",       // Caused by ticket expiration.
+	}
+
+	errMsg := strings.ToLower(err.Error())
+	for _, msg := range transientErrorMessages {
+		if strings.Contains(errMsg, msg) {
+			return true
+		}
+	}
+	return false
+}
+
 func prepareAIPayment(ctx context.Context, sess *AISession, outPixels int64) (worker.RequestEditorFn, *BalanceUpdate, error) {
 	// genSegCreds expects a stream.HLSSegment so in order to reuse it here we pass a dummy object
 	segCreds, err := genSegCreds(sess.BroadcastSession, &stream.HLSSegment{}, nil, false)

server/ai_process_test.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"errors"
 	"reflect"
 	"testing"
 
@@ -123,3 +124,35 @@ func TestEncodeReqMetadata(t *testing.T) {
 		})
 	}
 }
+
+func Test_isRetryableError(t *testing.T) {
+	tests := []struct {
+		name string
+		err  error
+		want bool
+	}{
+		{
+			name: "insufficient capacity error",
+			err:  errors.New("Insufficient capacity"),
+			want: true,
+		},
+		{
+			name: "INSUFFICIENT capacity ERROR",
+			err:  errors.New("Insufficient capacity"),
+			want: true,
+		},
+		{
+			name: "non-retryable error",
+			err:  errors.New("some other error"),
+			want: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := isRetryableError(tt.err); got != tt.want {
+				t.Errorf("isRetryableError() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}

server/ai_session.go

@@ -25,21 +25,23 @@ type AISession struct {
 }
 
 type AISessionPool struct {
-	selector  BroadcastSessionsSelector
-	sessMap   map[string]*BroadcastSession
-	inUseSess []*BroadcastSession
-	suspender *suspender
-	penalty   int
-	mu        sync.RWMutex
+	selector       BroadcastSessionsSelector
+	sessMap        map[string]*BroadcastSession
+	inUseSess      []*BroadcastSession
+	sessionsOnHold map[string][]*BroadcastSession
+	suspender      *suspender
+	penalty        int
+	mu             sync.RWMutex
 }
 
 func NewAISessionPool(selector BroadcastSessionsSelector, suspender *suspender, penalty int) *AISessionPool {
 	return &AISessionPool{
-		selector:  selector,
-		sessMap:   make(map[string]*BroadcastSession),
-		suspender: suspender,
-		penalty:   penalty,
-		mu:        sync.RWMutex{},
+		selector:       selector,
+		sessMap:        make(map[string]*BroadcastSession),
+		sessionsOnHold: make(map[string][]*BroadcastSession),
+		suspender:      suspender,
+		penalty:        penalty,
+		mu:             sync.RWMutex{},
 	}
 }
 
@@ -135,6 +137,31 @@ func (pool *AISessionPool) Remove(sess *BroadcastSession) {
 	pool.suspender.suspend(sess.Transcoder(), penalty)
 }
 
+func (pool *AISessionPool) PutSessionOnHold(id string, sess *BroadcastSession) {
+	pool.mu.Lock()
+	defer pool.mu.Unlock()
+
+	var sessionList []*BroadcastSession
+	holdList, ok := pool.sessionsOnHold[id]
+	if ok {
+		sessionList = holdList
+	}
+	pool.sessionsOnHold[id] = append(sessionList, sess)
+
+	delete(pool.sessMap, sess.Transcoder())
+	pool.inUseSess = removeSessionFromList(pool.inUseSess, sess)
+}
+
+func (pool *AISessionPool) ReleaseSessionsOnHold(id string) {
+	sessions, ok := pool.sessionsOnHold[id]
+	if !ok {
+		return
+	}
+
+	pool.Add(sessions)
+	delete(pool.sessionsOnHold, id)
+}
+
 func (pool *AISessionPool) Size() int {
 	pool.mu.RLock()
 	defer pool.mu.RUnlock()
@@ -245,7 +272,8 @@ func (sel *AISessionSelector) Select(ctx context.Context) *AISession {
 
 		// Refresh if the # of sessions across warm and cold pools falls below the smaller of the maxRefreshSessionsThreshold and
 		// 1/2 the total # of orchs that can be queried during discovery
-		if sel.warmPool.Size()+sel.coldPool.Size() < int(math.Min(maxRefreshSessionsThreshold, math.Ceil(float64(discoveryPoolSize)/2.0))) {
+		size := sel.warmPool.Size() + sel.coldPool.Size()
+		if size < int(math.Min(maxRefreshSessionsThreshold, math.Ceil(float64(discoveryPoolSize)/2.0))) {
 			return true
 		}
 
@@ -293,6 +321,30 @@ func (sel *AISessionSelector) Remove(sess *AISession) {
 	}
 }
 
+func (sel *AISessionSelector) PutSessionOnHold(id string, sess *AISession) {
+	// start cleanup func on first session on hold for request
+	_, warmPoolCleanupStarted := sel.warmPool.sessionsOnHold[id]
+	_, coldPoolCleanupStarted := sel.coldPool.sessionsOnHold[id]
+	if !warmPoolCleanupStarted && !coldPoolCleanupStarted {
+		go func(id string, sel *AISessionSelector) {
+			time.Sleep(sel.node.AIProcesssingRetryTimeout)
+			sel.ReleaseSessionsOnHold(id)
+		}(id, sel)
+	}
+
+	// put session on hold
+	if sess.Warm {
+		sel.warmPool.PutSessionOnHold(id, sess.BroadcastSession)
+	} else {
+		sel.coldPool.PutSessionOnHold(id, sess.BroadcastSession)
+	}
+}
+
+func (sel *AISessionSelector) ReleaseSessionsOnHold(id string) {
+	sel.warmPool.ReleaseSessionsOnHold(id)
+	sel.coldPool.ReleaseSessionsOnHold(id)
+}
+
 func (sel *AISessionSelector) Refresh(ctx context.Context) error {
 	// If we try to add new sessions to the pool the suspender
 	// should treat this as a refresh
@@ -421,6 +473,17 @@ func (c *AISessionManager) Remove(ctx context.Context, sess *AISession) error {
 	return nil
 }
 
+func (c *AISessionManager) PutSessionOnHold(ctx context.Context, id string, sess *AISession) error {
+	sel, err := c.getSelector(ctx, sess.Cap, sess.ModelID)
+	if err != nil {
+		return err
+	}
+
+	sel.PutSessionOnHold(id, sess)
+
+	return nil
+}
+
 func (c *AISessionManager) Complete(ctx context.Context, sess *AISession) error {
 	sel, err := c.getSelector(ctx, sess.Cap, sess.ModelID)
 	if err != nil {