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 @@ type LivepeerConfig struct {
 	OrchMinLivepeerVersion     *string
 	TestOrchAvail              *bool
 	AIRunnerImage              *string
+	AIRunnerImageOverrides     *string
 	KafkaBootstrapServers      *string
 	KafkaUsername              *string
 	KafkaPassword              *string
@@ -214,6 +215,7 @@ func DefaultLivepeerConfig() LivepeerConfig {
 	defaultAIModels := ""
 	defaultAIModelsDir := ""
 	defaultAIRunnerImage := "livepeer/ai-runner:latest"
+	defaultAIRunnerImageOverrides := ""
 	defaultLiveAIAuthWebhookURL := ""
 	defaultLivePaymentInterval := 5 * time.Second
 	defaultGatewayHost := ""
@@ -318,14 +320,15 @@ func DefaultLivepeerConfig() LivepeerConfig {
 		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 @@ func StartLivepeer(ctx context.Context, cfg LivepeerConfig) {
 			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

go.mod

@@ -15,7 +15,7 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/jaypipes/ghw v0.10.0
 	github.com/jaypipes/pcidb v1.0.0
-	github.com/livepeer/ai-worker v0.13.2
+	github.com/livepeer/ai-worker v0.13.3
 	github.com/livepeer/go-tools v0.3.6-0.20240130205227-92479de8531b
 	github.com/livepeer/livepeer-data v0.7.5-0.20231004073737-06f1f383fb18
 	github.com/livepeer/lpms v0.0.0-20250118014304-79e6dcf08057

go.sum

@@ -605,8 +605,8 @@ github.com/libp2p/go-netroute v0.2.0 h1:0FpsbsvuSnAhXFnCY0VLFbJOzaK0VnP0r1QT/o4n
 github.com/libp2p/go-netroute v0.2.0/go.mod h1:Vio7LTzZ+6hoT4CMZi5/6CpY3Snzh2vgZhWgxMNwlQI=
 github.com/libp2p/go-openssl v0.1.0 h1:LBkKEcUv6vtZIQLVTegAil8jbNpJErQ9AnT+bWV+Ooo=
 github.com/libp2p/go-openssl v0.1.0/go.mod h1:OiOxwPpL3n4xlenjx2h7AwSGaFSC/KZvf6gNdOBQMtc=
-github.com/livepeer/ai-worker v0.13.2 h1:rJs2utzibdYWKiN23QE1oJ+ZEB2NGus6BTcvrw7dAcM=
-github.com/livepeer/ai-worker v0.13.2/go.mod h1:rbcoIzQewbf5rvosCvG2M9DLg/ZMl7yMsaSce3svXFA=
+github.com/livepeer/ai-worker v0.13.3 h1:vcKUK56GRwiHIhz0UbNeKffFBAPtJPWuZNmcJrhAV8o=
+github.com/livepeer/ai-worker v0.13.3/go.mod h1:rbcoIzQewbf5rvosCvG2M9DLg/ZMl7yMsaSce3svXFA=
 github.com/livepeer/go-tools v0.3.6-0.20240130205227-92479de8531b h1:VQcnrqtCA2UROp7q8ljkh2XA/u0KRgVv0S1xoUvOweE=
 github.com/livepeer/go-tools v0.3.6-0.20240130205227-92479de8531b/go.mod h1:hwJ5DKhl+pTanFWl+EUpw1H7ukPO/H+MFpgA7jjshzw=
 github.com/livepeer/joy4 v0.1.2-0.20191121080656-b2fea45cbded h1:ZQlvR5RB4nfT+cOQee+WqmaDOgGtP2oDMhcVvR4L0yA=