Allow setting controller image version via env

During testing, or as an occasional workaround, it would be useful to
specify a model-runner controller image version. So, add the ability to
supply it via an environment variable.

Signed-off-by: Piotr Stankiewicz <[email protected]>

Author: Piotr Stankiewicz

pkg/standalone/containers.go

@@ -71,9 +71,9 @@ func CreateControllerContainer(ctx context.Context, dockerClient *client.Client,
 	var imageName string
 	switch gpu {
 	case gpupkg.GPUSupportCUDA:
-		imageName = ControllerImage + ":" + controllerImageTagCUDA
+		imageName = ControllerImage + ":" + controllerImageTagCUDA()
 	default:
-		imageName = ControllerImage + ":" + controllerImageTagCPU
+		imageName = ControllerImage + ":" + controllerImageTagCPU()
 	}
 
 	// Set up the container configuration.

pkg/standalone/images.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"os"
 
 	"github.com/docker/docker/api/types/image"
 	"github.com/docker/docker/client"
@@ -15,23 +16,37 @@ import (
 const (
 	// ControllerImage is the image used for the controller container.
 	ControllerImage = "docker/model-runner"
-	// controllerImageTagCPU is the image tag used for the controller container
+	// defaultControllerImageTagCPU is the image tag used for the controller container
 	// when running with the CPU backend.
-	controllerImageTagCPU = "latest"
-	// controllerImageTagCUDA is the image tag used for the controller container
+	defaultControllerImageTagCPU = "latest"
+	// defaultControllerImageTagCUDA is the image tag used for the controller container
 	// when running with the CUDA GPU backend.
-	controllerImageTagCUDA = "latest-cuda"
+	defaultControllerImageTagCUDA = "latest-cuda"
 )
 
+func controllerImageTagCPU() string {
+	if version, ok := os.LookupEnv("MODEL_RUNNER_CONTROLLER_VERSION"); ok {
+		return version
+	}
+	return defaultControllerImageTagCPU
+}
+
+func controllerImageTagCUDA() string {
+	if version, ok := os.LookupEnv("MODEL_RUNNER_CONTROLLER_VERSION"); ok {
+		return version + "-cuda"
+	}
+	return defaultControllerImageTagCUDA
+}
+
 // EnsureControllerImage ensures that the controller container image is pulled.
 func EnsureControllerImage(ctx context.Context, dockerClient *client.Client, gpu gpupkg.GPUSupport, printer StatusPrinter) error {
 	// Determine the target image.
 	var imageName string
 	switch gpu {
 	case gpupkg.GPUSupportCUDA:
-		imageName = ControllerImage + ":" + controllerImageTagCUDA
+		imageName = ControllerImage + ":" + controllerImageTagCUDA()
 	default:
-		imageName = ControllerImage + ":" + controllerImageTagCPU
+		imageName = ControllerImage + ":" + controllerImageTagCPU()
 	}
 
 	// Perform the pull.
@@ -65,13 +80,13 @@ func EnsureControllerImage(ctx context.Context, dockerClient *client.Client, gpu
 // PruneControllerImages removes any unused controller container images.
 func PruneControllerImages(ctx context.Context, dockerClient *client.Client, printer StatusPrinter) error {
 	// Remove the standard image, if present.
-	imageNameCPU := ControllerImage + ":" + controllerImageTagCPU
+	imageNameCPU := ControllerImage + ":" + controllerImageTagCPU()
 	if _, err := dockerClient.ImageRemove(ctx, imageNameCPU, image.RemoveOptions{}); err == nil {
 		printer.Println("Removed image", imageNameCPU)
 	}
 
 	// Remove the CUDA GPU image, if present.
-	imageNameCUDA := ControllerImage + ":" + controllerImageTagCUDA
+	imageNameCUDA := ControllerImage + ":" + controllerImageTagCUDA()
 	if _, err := dockerClient.ImageRemove(ctx, imageNameCUDA, image.RemoveOptions{}); err == nil {
 		printer.Println("Removed image", imageNameCUDA)
 	}