Enhance aiRunnerImage flag

cmd/livepeer/livepeer.go

@@ -161,7 +161,8 @@ func parseLivepeerConfig() starter.LivepeerConfig {
 	cfg.AIWorker = flag.Bool("aiWorker", *cfg.AIWorker, "Set to true to run an AI worker")
 	cfg.AIModels = flag.String("aiModels", *cfg.AIModels, "Set models (pipeline:model_id) for AI worker to load upon initialization")
 	cfg.AIModelsDir = flag.String("aiModelsDir", *cfg.AIModelsDir, "Set directory where AI model weights are stored")
-	cfg.AIRunnerImage = flag.String("aiRunnerImage", *cfg.AIRunnerImage, "Set the docker image for the AI runner: Example - livepeer/ai-runner:0.0.1")
+	cfg.AIRunnerImage = flag.String("aiRunnerImage", *cfg.AIRunnerImage, "[Deprecated] Specify the base Docker image for the AI runner. Example: livepeer/ai-runner:0.0.1. Use -aiRunnerImageOverrides instead.")
+	cfg.AIRunnerImageOverrides = flag.String("aiRunnerImageOverrides", *cfg.AIRunnerImageOverrides, `Specify overrides for the Docker images used by the AI runner. Example: '{"default": "livepeer/ai-runner:v1.0", "batch": {"text-to-speech": "livepeer/ai-runner:text-to-speech-v1.0"}, "live": {"another-pipeline": "livepeer/ai-runner:another-pipeline-v1.0"}}'`)
 
 	// Live AI:
 	cfg.MediaMTXApiPassword = flag.String("mediaMTXApiPassword", "", "HTTP basic auth password for MediaMTX API requests")

cmd/livepeer/starter/starter.go

@@ -164,6 +164,7 @@
 	OrchMinLivepeerVersion     *string
 	TestOrchAvail              *bool
 	AIRunnerImage              *string
+	AIRunnerImageOverrides     *string
 	KafkaBootstrapServers      *string
 	KafkaUsername              *string
 	KafkaPassword              *string
@@ -214,6 +215,7 @@
 	defaultAIModels := ""
 	defaultAIModelsDir := ""
 	defaultAIRunnerImage := "livepeer/ai-runner:latest"
+	defaultAIRunnerImageOverrides := ""
 	defaultLiveAIAuthWebhookURL := ""
 	defaultLivePaymentInterval := 5 * time.Second
 	defaultGatewayHost := ""
@@ -318,14 +320,15 @@
 		TestTranscoder:       &defaultTestTranscoder,
 
 		// AI:
-		AIServiceRegistry:    &defaultAIServiceRegistry,
-		AIWorker:             &defaultAIWorker,
-		AIModels:             &defaultAIModels,
-		AIModelsDir:          &defaultAIModelsDir,
-		AIRunnerImage:        &defaultAIRunnerImage,
-		LiveAIAuthWebhookURL: &defaultLiveAIAuthWebhookURL,
-		LivePaymentInterval:  &defaultLivePaymentInterval,
-		GatewayHost:          &defaultGatewayHost,
+		AIServiceRegistry:      &defaultAIServiceRegistry,
+		AIWorker:               &defaultAIWorker,
+		AIModels:               &defaultAIModels,
+		AIModelsDir:            &defaultAIModelsDir,
+		AIRunnerImage:          &defaultAIRunnerImage,
+		AIRunnerImageOverrides: &defaultAIRunnerImageOverrides,
+		LiveAIAuthWebhookURL:   &defaultLiveAIAuthWebhookURL,
+		LivePaymentInterval:    &defaultLivePaymentInterval,
+		GatewayHost:            &defaultGatewayHost,
 
 		// Onchain:
 		EthAcctAddr:             &defaultEthAcctAddr,
@@ -1211,7 +1214,24 @@
 			return
 		}
 
-		n.AIWorker, err = worker.NewWorker(*cfg.AIRunnerImage, gpus, modelsDir)
+		// Retrieve image overrides from the config.
+		var imageOverrides worker.ImageOverrides
+		if *cfg.AIRunnerImageOverrides != "" {
+			if err := json.Unmarshal([]byte(*cfg.AIRunnerImageOverrides), &imageOverrides); err != nil {
+				glog.Errorf("Error unmarshaling image overrides: %v", err)
+				return
+			}
+		}
+
+		// Backwards compatibility for deprecated flags.
+		if *cfg.AIRunnerImage != "" {
+			glog.Warning("-aiRunnerImage flag is deprecated and will be removed in a future release. Please use -aiWorkerImageOverrides instead")
+			if imageOverrides.Default == "" {
+				imageOverrides.Default = *cfg.AIRunnerImage
+			}
+		}
+
+		n.AIWorker, err = worker.NewWorker(imageOverrides, gpus, modelsDir)
 		if err != nil {
 			glog.Errorf("Error starting AI worker: %v", err)
 			return