fix(ai): fix combined orchestrator capacity management

core/ai_orchestrator.go

@@ -398,18 +398,44 @@
 }
 
 // CheckAICapacity verifies if the orchestrator can process a request for a specific pipeline and modelID.
-func (orch *orchestrator) CheckAICapacity(pipeline, modelID string) bool {
+func (orch *orchestrator) CheckAICapacity(pipeline, modelID string) (bool, chan<- bool) {
+	var hasCapacity bool
 	if orch.node.AIWorker != nil {
 		// confirm local worker has capacity
-		return orch.node.AIWorker.HasCapacity(pipeline, modelID)
+		if pipeline == "live-video-to-video" {
+			return orch.node.AIWorker.HasCapacity(pipeline, modelID), nil
+		}
+
+		// batch pipelines manage the capacity at the Orchestrator level to manage local ai-worker capacity
+		err := orch.node.ReserveAICapability(pipeline, modelID)
+		if err == nil {
+			hasCapacity = true
+		}
 	} else {
 		// remote workers: RemoteAIWorkerManager only selects remote workers if they have capacity for the pipeline/model
+		// live-video-to-video is not using remote workers currently
 		if orch.node.AIWorkerManager != nil {
-			return orch.node.AIWorkerManager.workerHasCapacity(pipeline, modelID)
-		} else {
-			return false
+			hasCapacity = orch.node.AIWorkerManager.workerHasCapacity(pipeline, modelID)
 		}
 	}
+
+	if !hasCapacity {
+		return false, nil
+	}
+
+	// reserve AI capacity for the pipeline and modelID
+	releaseCapacity := make(chan bool)
+
+	go func() {
+		<-releaseCapacity
+		orch.node.ReleaseAICapability(pipeline, modelID)
+		glog.Infof("Released AI capacity for pipeline=%s model_id=%s", pipeline, modelID)
+		close(releaseCapacity)
+
+	}()
+
+	return true, releaseCapacity
+
 }
 
 func (orch *orchestrator) GetLiveAICapacity() worker.Capacity {
@@ -545,6 +571,7 @@
 	// local AIWorker processes job if combined orchestrator/ai worker
 	if orch.node.AIWorker != nil {
 		workerResp, err := orch.node.TextToImage(ctx, req)
+
 		if err == nil {
 			return orch.node.saveLocalAIWorkerResults(ctx, *workerResp, requestID, "image/png")
 		} else {
@@ -578,6 +605,7 @@
 	// local AIWorker processes job if combined orchestrator/ai worker
 	if orch.node.AIWorker != nil {
 		workerResp, err := orch.node.LiveVideoToVideo(ctx, req)
+
 		if err == nil {
 			return orch.node.saveLocalAIWorkerResults(ctx, *workerResp, requestID, "application/json")
 		} else {
@@ -611,6 +639,7 @@
 	// local AIWorker processes job if combined orchestrator/ai worker
 	if orch.node.AIWorker != nil {
 		workerResp, err := orch.node.ImageToImage(ctx, req)
+
 		if err == nil {
 			return orch.node.saveLocalAIWorkerResults(ctx, *workerResp, requestID, "image/png")
 		} else {
@@ -655,6 +684,7 @@
 	// local AIWorker processes job if combined orchestrator/ai worker
 	if orch.node.AIWorker != nil {
 		workerResp, err := orch.node.ImageToVideo(ctx, req)
+
 		if err == nil {
 			return orch.node.saveLocalAIWorkerResults(ctx, *workerResp, requestID, "video/mp4")
 		} else {
@@ -699,6 +729,7 @@
 	// local AIWorker processes job if combined orchestrator/ai worker
 	if orch.node.AIWorker != nil {
 		workerResp, err := orch.node.Upscale(ctx, req)
+
 		if err == nil {
 			return orch.node.saveLocalAIWorkerResults(ctx, *workerResp, requestID, "image/png")
 		} else {
@@ -880,6 +911,7 @@
 	// local AIWorker processes job if combined orchestrator/ai worker
 	if orch.node.AIWorker != nil {
 		workerResp, err := orch.node.TextToSpeech(ctx, req)
+
 		if err == nil {
 			return orch.node.saveLocalAIWorkerResults(ctx, *workerResp, requestID, "audio/wav")
 		} else {

core/ai_test.go

@@ -515,8 +515,9 @@ func TestCheckAICapacity(t *testing.T) {
 	n.Capabilities = createAIWorkerCapabilities()
 	n.AIWorker = &wkr
 	// Test when local AI worker has capacity
-	hasCapacity := o.CheckAICapacity("text-to-image", "livepeer/model1")
+	hasCapacity, releaseCapacity := o.CheckAICapacity("text-to-image", "livepeer/model1")
 	assert.True(t, hasCapacity)
+	releaseCapacity <- true
 
 	o.node.AIWorker = nil
 	o.node.AIWorkerManager = NewRemoteAIWorkerManager()
@@ -534,12 +535,15 @@ func TestCheckAICapacity(t *testing.T) {
 	}()
 	time.Sleep(1 * time.Millisecond) // allow the workers to activate
 
-	hasCapacity = o.CheckAICapacity("text-to-image", "livepeer/model1")
+	hasCapacity, releaseCapacity = o.CheckAICapacity("text-to-image", "livepeer/model1")
 	assert.True(t, hasCapacity)
+	assert.NotNil(t, releaseCapacity)
+	releaseCapacity <- true
 
 	// Test when remote AI worker does not have capacity
-	hasCapacity = o.CheckAICapacity("text-to-image", "livepeer/model2")
+	hasCapacity, releaseCapacity = o.CheckAICapacity("text-to-image", "livepeer/model2")
 	assert.False(t, hasCapacity)
+	assert.Nil(t, releaseCapacity)
 }
 func TestRemoteAIWorkerProcessPipelines(t *testing.T) {
 	drivers.NodeStorage = drivers.NewMemoryDriver(nil)

server/ai_http.go

@@ -131,8 +131,10 @@
 		}
 
 		// Check if there is capacity for the request
-		if !orch.CheckAICapacity(pipeline, modelID) {
-			respondWithError(w, fmt.Sprintf("Insufficient capacity for pipeline=%v modelID=%v", pipeline, modelID), http.StatusServiceUnavailable)
+		hasCapacity, _ := orch.CheckAICapacity(pipeline, modelID)
+		if !hasCapacity {
+			clog.Errorf(ctx, "Insufficient capacity for pipeline=%v modelID=%v", pipeline, modelID)
+			respondWithError(w, "insufficient capacity", http.StatusServiceUnavailable)
 			return
 		}
 
@@ -497,8 +499,11 @@
 	manifestID := core.ManifestID(strconv.Itoa(int(cap)) + "_" + modelID)
 
 	// Check if there is capacity for the request.
-	if !orch.CheckAICapacity(pipeline, modelID) {
-		respondWithError(w, fmt.Sprintf("Insufficient capacity for pipeline=%v modelID=%v", pipeline, modelID), http.StatusServiceUnavailable)
+	// Capability capacity is reserved if available and released when response is received
+	hasCapacity, releaseCapacity := orch.CheckAICapacity(pipeline, modelID)
+	if !hasCapacity {
+		clog.Errorf(ctx, "Insufficient capacity for pipeline=%v modelID=%v", pipeline, modelID)
+		respondWithError(w, "insufficient capacity", http.StatusServiceUnavailable)
 		return
 	}
 
@@ -528,6 +533,7 @@
 
 	start := time.Now()
 	resp, err := submitFn(ctx)
+
 	if err != nil {
 		if monitor.Enabled {
 			monitor.AIProcessingError(err.Error(), pipeline, modelID, sender.Hex())
@@ -614,6 +620,7 @@
 		flusher, ok := w.(http.Flusher)
 		if !ok {
 			http.Error(w, "Streaming unsupported!", http.StatusInternalServerError)
+			releaseCapacity <- true
 			return
 		}
 
@@ -631,8 +638,12 @@
 				break
 			}
 		}
+		//release capacity after streaming is done
+		releaseCapacity <- true
+
 	} else {
 		// Non-streaming response
+		releaseCapacity <- true
 		w.Header().Set("Content-Type", "application/json")
 		w.WriteHeader(http.StatusOK)
 		_ = json.NewEncoder(w).Encode(resp)

server/rpc.go

@@ -52,7 +52,7 @@
 	Sign([]byte) ([]byte, error)
 	VerifySig(ethcommon.Address, string, []byte) bool
 	CheckCapacity(core.ManifestID) error
-	CheckAICapacity(pipeline, modelID string) bool
+	CheckAICapacity(pipeline, modelID string) (bool, chan<- bool)
 	GetLiveAICapacity() worker.Capacity
 	TranscodeSeg(context.Context, *core.SegTranscodingMetadata, *stream.HLSSegment) (*core.TranscodeResult, error)
 	ServeTranscoder(stream net.Transcoder_RegisterTranscoderServer, capacity int, capabilities *net.Capabilities)
@@ -390,7 +390,8 @@
 	if liveCap, ok := caps.Constraints.PerCapability[uint32(core.Capability_LiveVideoToVideo)]; ok {
 		pipeline := "live-video-to-video"
 		for modelID := range liveCap.GetModels() {
-			if orch.CheckAICapacity(pipeline, modelID) {
+			hasCapacity, _ := orch.CheckAICapacity(pipeline, modelID)
+			if hasCapacity {
 				// It has capacity for at least one of the requested models
 				return nil
 			}

server/rpc_test.go

@@ -235,8 +235,8 @@ func (r *stubOrchestrator) LiveVideoToVideo(ctx context.Context, requestID strin
 	return nil, nil
 }
 
-func (r *stubOrchestrator) CheckAICapacity(pipeline, modelID string) bool {
-	return true
+func (r *stubOrchestrator) CheckAICapacity(pipeline, modelID string) (bool, chan<- bool) {
+	return true, nil
 }
 func (r *stubOrchestrator) AIResults(job int64, res *core.RemoteAIWorkerResult) {
 }
@@ -1486,8 +1486,8 @@ func (r *mockOrchestrator) TextToSpeech(ctx context.Context, requestID string, r
 func (r *mockOrchestrator) LiveVideoToVideo(ctx context.Context, requestID string, req worker.GenLiveVideoToVideoJSONRequestBody) (interface{}, error) {
 	return nil, nil
 }
-func (r *mockOrchestrator) CheckAICapacity(pipeline, modelID string) bool {
-	return true
+func (r *mockOrchestrator) CheckAICapacity(pipeline, modelID string) (bool, chan<- bool) {
+	return true, nil
 }
 func (r *mockOrchestrator) AIResults(job int64, res *core.RemoteAIWorkerResult) {