feat(ai): enable worker image pulling (#3279)

This commit ensures that the docker containers the worker uses are
pulled during startup. They use the new changes implemented in
livepeer/ai-runner#200.

CHANGELOG_PENDING.md

@@ -2,6 +2,8 @@
 
 ## v0.X.X
 
+- [#3279](https://github.com/livepeer/go-livepeer/pull/3279) - Enable automatic worker image pulling.
+
 ### Breaking Changes 🚨🚨
 
 ### Features ⚒

cmd/livepeer/starter/starter.go

@@ -1242,6 +1242,12 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
 					modelConstraint.Capacity = config.Capacity
 				}
 
+				// Ensure the AI worker has the image needed to serve the job.
+				err := n.AIWorker.EnsureImageAvailable(ctx, config.Pipeline, config.ModelID)
+				if err != nil {
+					glog.Errorf("Error ensuring AI worker image available for %v: %v", config.Pipeline, err)
+				}
+
 				if config.Warm || config.URL != "" {
 					// Register external container endpoint if URL is provided.
 					endpoint := worker.RunnerEndpoint{URL: config.URL, Token: config.Token}

core/ai.go

@@ -30,7 +30,8 @@ type AI interface {
 	LiveVideoToVideo(context.Context, worker.GenLiveVideoToVideoJSONRequestBody) (*worker.LiveVideoToVideoResponse, error)
 	Warm(context.Context, string, string, worker.RunnerEndpoint, worker.OptimizationFlags) error
 	Stop(context.Context) error
-	HasCapacity(pipeline, modelID string) bool
+	HasCapacity(string, string) bool
+	EnsureImageAvailable(context.Context, string, string) error
 }
 
 // Custom type to parse a big.Rat from a JSON number.

core/ai_test.go

@@ -679,6 +679,10 @@ func (a *stubAIWorker) HasCapacity(pipeline, modelID string) bool {
 	return true
 }
 
+func (a *stubAIWorker) EnsureImageAvailable(ctx context.Context, pipeline string, modelID string) error {
+	return nil
+}
+
 type StubAIWorkerServer struct {
 	manager      *RemoteAIWorkerManager
 	SendError    error

server/ai_worker_test.go

@@ -649,3 +649,8 @@ func (a *stubAIWorker) HasCapacity(pipeline, modelID string) bool {
 	a.Called++
 	return true
 }
+
+func (a *stubAIWorker) EnsureImageAvailable(ctx context.Context, pipeline string, modelID string) error {
+	a.Called++
+	return nil
+}