neonvm-controller: Add image mapping feature

go.mod

@@ -224,3 +224,5 @@ require (
 	sigs.k8s.io/structured-merge-diff/v4 v4.4.1 // indirect
 	sigs.k8s.io/yaml v1.4.0 // indirect
 )
+
+require github.com/BurntSushi/toml v1.4.0

go.sum

@@ -18,6 +18,8 @@ github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1 h1:UQHMgLO+TxOEl
 github.com/Azure/go-ansiterm v0.0.0-20210617225240-d185dfc1b5a1/go.mod h1:xomTg63KZ2rFqZQzSB4Vz2SUXa1BpHTVz9L5PTmPC4E=
 github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2 h1:XHOnouVk1mxXfQidrMEnLlPk9UMeRtyBTnEFtxkV0kU=
 github.com/AzureAD/microsoft-authentication-library-for-go v1.2.2/go.mod h1:wP83P5OoQ5p6ip3ScPr0BAq0BvuPAvacpEuSzyouqAI=
+github.com/BurntSushi/toml v1.4.0 h1:kuoIxZQy2WRRk1pttg9asf+WVv6tWQuBNVmK8+nqPr0=
+github.com/BurntSushi/toml v1.4.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho=
 github.com/Microsoft/go-winio v0.6.0 h1:slsWYD/zyx7lCXoZVlvQrj0hPTM1HI4+v1sIda2yDvg=
 github.com/Microsoft/go-winio v0.6.0/go.mod h1:cTAf44im0RAYeL23bpB+fzCyDH2MJiz2BO69KH/soAE=
 github.com/NYTimes/gziphandler v1.1.1 h1:ZUDjpQae29j0ryrS0u/B8HZfJBtBQHjqw2rQ2cqUQ3I=

neonvm-controller/cmd/main.go

@@ -101,6 +101,7 @@ func main() {
 	var failurePendingPeriod time.Duration
 	var failingRefreshInterval time.Duration
 	var atMostOnePod bool
+	var imageMapPath string
 	flag.StringVar(&metricsAddr, "metrics-bind-address", ":8080", "The address the metric endpoint binds to.")
 	flag.StringVar(&probeAddr, "health-probe-bind-address", ":8081", "The address the probe endpoint binds to.")
 	flag.BoolVar(&enableLeaderElection, "leader-elect", false,
@@ -144,6 +145,7 @@ func main() {
 	flag.BoolVar(&atMostOnePod, "at-most-one-pod", false,
 		"If true, the controller will ensure that at most one pod is running at a time. "+
 			"Otherwise, the outdated pod might be left to terminate, while the new one is already running.")
+	flag.StringVar(&imageMapPath, "imagemap", "", "Path to an image mappings file, for overriding images specified in the VirtualMachine spec")
 	flag.Parse()
 
 	logConfig := zap.NewProductionConfig()
@@ -195,6 +197,7 @@ func main() {
 		FailingRefreshInterval:  failingRefreshInterval,
 		AtMostOnePod:            atMostOnePod,
 		DefaultCPUScalingMode:   defaultCpuScalingMode,
+		ImageMapPath:            imageMapPath,
 	}
 
 	vmReconciler := &controllers.VMReconciler{

pkg/neonvm/controllers/config.go

@@ -51,4 +51,9 @@ type ReconcilerConfig struct {
 	AtMostOnePod bool
 	// DefaultCPUScalingMode is the default CPU scaling mode that will be used for VMs with empty spec.cpuScalingMode
 	DefaultCPUScalingMode vmv1.CpuScalingMode
+
+	// Path to an image mappings file, for overriding images specified in an VirtualMachine spec.
+	// This is meant for local development, to allow rapidly changing images without changing the
+	// component that creates the VirtualMachine object.
+	ImageMapPath string
 }

pkg/neonvm/controllers/functests/vm_controller_test.go

@@ -115,6 +115,7 @@ var _ = Describe("VirtualMachine controller", func() {
 					FailingRefreshInterval:  1 * time.Minute,
 					AtMostOnePod:            false,
 					DefaultCPUScalingMode:   vmv1.CpuScalingModeQMP,
+					ImageMapPath:            "",
 				},
 			}
 

pkg/neonvm/controllers/image_map.go

@@ -0,0 +1,81 @@
+package controllers
+
+// An "image map" allows quick overriding of the images specified in VirtualMachine objects.
+// That's handy during local development, for example, so that you can quickly swap images
+// without having to change the compute image name in the control plane database. The image
+// map is used by the compute.Tiltfile in the cloud repository to allow Tilt to replace the
+// compute image; the Tiltfile commands regenerate the image mapping file with the tilt-generated
+// image names whenever the compute images are rebuilt.
+//
+// An image map file consists of pairs of image names, specifying that when the VirtualMachine
+// spec contains X, it is replaced with Y. If an image name is not present in the image map, it
+// is used without changes.  For example, to replace v17 and v16 images with local versions,
+// you could use this file:
+//
+// ```
+// "vm-compute-node-v17:latest" = "vm-compute-node-v17:dev"
+// "vm-compute-node-v16:latest" = "vm-compute-node-v16:dev"
+// ```
+//
+// We use a toml file parser to parse the file, so toml syntax rules on comments and escaping
+// apply.
+
+import (
+	"context"
+	"os"
+
+	"github.com/BurntSushi/toml"
+	"sigs.k8s.io/controller-runtime/pkg/log"
+)
+
+type ImageMap map[string]string
+
+func EmptyImageMap() ImageMap {
+	return map[string]string{}
+}
+
+func tryLoadImageMap(ctx context.Context, path string) ImageMap {
+	if path == "" {
+		return EmptyImageMap()
+	}
+
+	log := log.FromContext(ctx)
+
+	newMap, err := loadImageMap(path)
+	if err != nil {
+		log.Error(err, "could not read image mapping file",
+			"ImageMapPath", path)
+	}
+
+	content, err := os.ReadFile(path)
+	if err != nil {
+		log.Error(err, "xcould not read image mapping file",
+			"ImageMapPath", path)
+		return EmptyImageMap()
+	}
+
+	log.Info("imagemap loaded", "map", newMap, "content", content)
+	return newMap
+}
+
+func loadImageMap(path string) (ImageMap, error) {
+	content, err := os.ReadFile(path)
+	if err != nil {
+		return EmptyImageMap(), err
+	}
+	var newMap map[string]string
+	err = toml.Unmarshal(content, &newMap)
+	if err != nil {
+		return EmptyImageMap(), err
+	}
+	return newMap, nil
+}
+
+func (imageMap *ImageMap) mapImage(image string) string {
+	mappedImage := (*imageMap)[image]
+	if mappedImage != "" {
+		return mappedImage
+	} else {
+		return image
+	}
+}

pkg/neonvm/controllers/vm_controller.go

@@ -81,6 +81,8 @@ type VMReconciler struct {
 	Recorder record.EventRecorder
 	Config   *ReconcilerConfig
 
+	ImageMap ImageMap `exhaustruct:"optional"`
+
 	Metrics ReconcilerMetrics `exhaustruct:"optional"`
 }
 
@@ -435,8 +437,11 @@ func (r *VMReconciler) doReconcile(ctx context.Context, vm *vmv1.VirtualMachine)
 				}
 			}
 
+			// Reload the image map every time, so that we pick up any changes quickly.
+			imageMap := tryLoadImageMap(ctx, r.Config.ImageMapPath)
+
 			// Define a new pod
-			pod, err := r.podForVirtualMachine(vm, sshSecret)
+			pod, err := r.podForVirtualMachine(vm, sshSecret, imageMap)
 			if err != nil {
 				log.Error(err, "Failed to define new Pod resource for VirtualMachine")
 				return err
@@ -1047,8 +1052,9 @@ func extractVirtualMachineResourcesJSON(spec vmv1.VirtualMachineSpec) string {
 func (r *VMReconciler) podForVirtualMachine(
 	vm *vmv1.VirtualMachine,
 	sshSecret *corev1.Secret,
+	imageMap ImageMap,
 ) (*corev1.Pod, error) {
-	pod, err := podSpec(vm, sshSecret, r.Config)
+	pod, err := podSpec(vm, sshSecret, r.Config, imageMap)
 	if err != nil {
 		return nil, err
 	}
@@ -1182,6 +1188,7 @@ func podSpec(
 	vm *vmv1.VirtualMachine,
 	sshSecret *corev1.Secret,
 	config *ReconcilerConfig,
+	imageMap ImageMap,
 ) (*corev1.Pod, error) {
 	runnerVersion := api.RunnerProtoV1
 	labels := labelsForVirtualMachine(vm, &runnerVersion)
@@ -1193,6 +1200,7 @@ func podSpec(
 	if err != nil {
 		return nil, err
 	}
+	image = imageMap.mapImage(image)
 
 	vmSpecJson, err := json.Marshal(vm.Spec)
 	if err != nil {
@@ -1242,7 +1250,7 @@ func podSpec(
 			Affinity:                      affinity,
 			InitContainers: []corev1.Container{
 				{
-					Image:           vm.Spec.Guest.RootDisk.Image,
+					Image:           imageMap.mapImage(vm.Spec.Guest.RootDisk.Image),
 					Name:            "init",
 					ImagePullPolicy: vm.Spec.Guest.RootDisk.ImagePullPolicy,
 					VolumeMounts: []corev1.VolumeMount{{
@@ -1425,14 +1433,14 @@ func podSpec(
 
 	// If a custom neonvm-runner image is requested, use that instead:
 	if vm.Spec.RunnerImage != nil {
-		pod.Spec.Containers[0].Image = *vm.Spec.RunnerImage
+		pod.Spec.Containers[0].Image = imageMap.mapImage(*vm.Spec.RunnerImage)
 	}
 
 	// If a custom kernel is used, add that image:
 	if vm.Spec.Guest.KernelImage != nil {
 		pod.Spec.Containers[0].Args = append(pod.Spec.Containers[0].Args, "-kernelpath=/vm/images/vmlinuz")
 		pod.Spec.InitContainers = append(pod.Spec.InitContainers, corev1.Container{
-			Image:           *vm.Spec.Guest.KernelImage,
+			Image:           imageMap.mapImage(*vm.Spec.Guest.KernelImage),
 			Name:            "init-kernel",
 			ImagePullPolicy: vm.Spec.Guest.RootDisk.ImagePullPolicy,
 			Args:            []string{"cp", "/vmlinuz", "/vm/images/vmlinuz"},

pkg/neonvm/controllers/vm_controller_unit_test.go

@@ -124,6 +124,7 @@ func newTestParams(t *testing.T) *testParams {
 			FailingRefreshInterval:  time.Minute,
 			AtMostOnePod:            false,
 			DefaultCPUScalingMode:   vmv1.CpuScalingModeQMP,
+			ImageMapPath:            "",
 		},
 		Metrics: reconcilerMetrics,
 	}

pkg/neonvm/controllers/vmmigration_controller.go

@@ -254,8 +254,10 @@ func (r *VirtualMachineMigrationReconciler) Reconcile(ctx context.Context, req c
 				}
 			}
 
+			imageMap := tryLoadImageMap(ctx, r.Config.ImageMapPath)
+
 			// Define a new target pod
-			tpod, err := r.targetPodForVirtualMachine(vm, migration, sshSecret)
+			tpod, err := r.targetPodForVirtualMachine(vm, migration, sshSecret, imageMap)
 			if err != nil {
 				log.Error(err, "Failed to generate Target Pod spec")
 				return ctrl.Result{}, err
@@ -699,12 +701,13 @@ func (r *VirtualMachineMigrationReconciler) targetPodForVirtualMachine(
 	vm *vmv1.VirtualMachine,
 	migration *vmv1.VirtualMachineMigration,
 	sshSecret *corev1.Secret,
+	imageMap ImageMap,
 ) (*corev1.Pod, error) {
 	if err := vm.Spec.Guest.ValidateMemorySize(); err != nil {
 		return nil, fmt.Errorf("cannot create target pod because memory is invalid: %w", err)
 	}
 
-	pod, err := podSpec(vm, sshSecret, r.Config)
+	pod, err := podSpec(vm, sshSecret, r.Config, imageMap)
 	if err != nil {
 		return nil, err
 	}