modified vmi/vm reconcile to also update vm with correct device names…

… based on vgpu annotation when VM is stopped

pkg/controller/setup.go

@@ -26,7 +26,7 @@ import (
 	"github.com/harvester/pcidevices/pkg/controller/pcideviceclaim"
 	"github.com/harvester/pcidevices/pkg/controller/sriovdevice"
 	"github.com/harvester/pcidevices/pkg/controller/usbdevice"
-	"github.com/harvester/pcidevices/pkg/controller/virtualmachineinstance"
+	"github.com/harvester/pcidevices/pkg/controller/virtualmachine"
 	"github.com/harvester/pcidevices/pkg/crd"
 	ctldevices "github.com/harvester/pcidevices/pkg/generated/controllers/devices.harvesterhci.io"
 	ctlkubevirt "github.com/harvester/pcidevices/pkg/generated/controllers/kubevirt.io"
@@ -105,7 +105,7 @@ func Setup(ctx context.Context, cfg *rest.Config, _ *runtime.Scheme) error {
 		sriovdevice.Register,
 		nodecleanup.Register,
 		gpudevice.Register,
-		virtualmachineinstance.Register,
+		virtualmachine.Register,
 	}
 
 	for _, register := range registers {

pkg/controller/virtualmachineinstance/virtualmachineinstance.go → pkg/controller/virtualmachine/virtualmachine.go

@@ -1,4 +1,4 @@
-package virtualmachineinstance
+package virtualmachine
 
 import (
 	"bufio"
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"reflect"
 	"slices"
 	"strings"
 
@@ -64,6 +65,7 @@ func Register(ctx context.Context, management *config.FactoryManager) error {
 		nodeName:       nodeName,
 	}
 	vmi.OnChange(ctx, "virtual-machine-instance-handler", h.OnVMIChange)
+	vmi.OnRemove(ctx, "virtual-machine-deletion", h.OnVMIDeletion)
 	return nil
 }
 
@@ -115,7 +117,12 @@ func (h *Handler) trackDevices(vmi *kubevirtv1.VirtualMachineInstance) error {
 	if err != nil {
 		return err
 	}
+	return h.reconcileDeviceAllocationDetails(vmi, envMap)
+}
 
+// reconcileDeviceAllocationDetails will reconcile envMap into device allocation annotation on vmi
+// has been split into its own method to simplify testing
+func (h *Handler) reconcileDeviceAllocationDetails(vmi *kubevirtv1.VirtualMachineInstance, envMap map[string]string) error {
 	var pciDeviceMap, vGPUMap map[string]string
 	selector := map[string]string{
 		"nodename": vmi.Status.NodeName,
@@ -360,3 +367,109 @@ func (h *Handler) reconcileVMResourceAllocationAnnotation(vmi *kubevirtv1.Virtua
 	}
 	return err
 }
+
+// OnVMIDeletion will update the VM spec with actual device allocation details
+// this simplifies deletion of devices from the VM object without needing any changes
+// it needs to be done during the VM shutdown avoid the object generation to differ between VM and VMI object
+// and this avoids the generation warning being reported in the UI
+func (h *Handler) OnVMIDeletion(_ string, vmi *kubevirtv1.VirtualMachineInstance) (*kubevirtv1.VirtualMachineInstance, error) {
+	if vmi == nil || vmi.DeletionTimestamp == nil {
+		return vmi, nil
+	}
+
+	// no host or GPU devices present so nothing needed to be done
+	if len(vmi.Spec.Domain.Devices.GPUs) == 0 && len(vmi.Spec.Domain.Devices.HostDevices) == 0 {
+		logrus.WithFields(logrus.Fields{
+			"name":      vmi.Name,
+			"namespace": vmi.Namespace,
+			"hostname":  vmi.Status.NodeName,
+		}).Debug("skipping vmi as it has no hostdevices or GPUs")
+		return vmi, nil
+	}
+
+	vmObj, err := h.vmCache.Get(vmi.Namespace, vmi.Name)
+	if err != nil {
+		return vmi, fmt.Errorf("error fetching vm object for vmi %s/%s: %v", vmi.Namespace, vmi.Name, err)
+	}
+
+	// no annotations so nothing to do
+	if vmObj.Annotations == nil {
+		logrus.WithFields(logrus.Fields{
+			"name":      vmi.Name,
+			"namespace": vmi.Namespace,
+			"hostname":  vmi.Status.NodeName,
+		}).Debug("skipping vmi as it has no annotations defined")
+		return vmi, nil
+	}
+
+	if vmi.Status.NodeName != h.nodeName {
+		logrus.WithFields(logrus.Fields{
+			"name":      vmi.Name,
+			"namespace": vmi.Namespace,
+			"hostname":  vmi.Status.NodeName,
+		}).Debug("skipping vmi as it is not scheduled on current node")
+		return vmi, nil
+	}
+
+	val, ok := vmObj.Annotations[v1beta1.DeviceAllocationKey]
+	if !ok {
+		// no device allocation annotations, nothing to do
+		logrus.WithFields(logrus.Fields{
+			"name":      vmi.Name,
+			"namespace": vmi.Namespace,
+			"hostname":  vmi.Status.NodeName,
+		}).Debug("skipping vmi as it has no device allocation annotation")
+		return vmi, nil
+	}
+
+	allocationDetails := &v1beta1.AllocationDetails{}
+	err = json.Unmarshal([]byte(val), allocationDetails)
+	if err != nil {
+		return vmi, fmt.Errorf("error unmarshalling allocation details annotation: %v", err)
+	}
+
+	vmObjCopy := vmObj.DeepCopy()
+	if len(vmi.Spec.Domain.Devices.HostDevices) > 0 {
+		patchHostDevices(vmObj, allocationDetails.HostDevices)
+	}
+
+	if len(vmi.Spec.Domain.Devices.GPUs) > 0 {
+		patchGPUDevices(vmObj, allocationDetails.GPUs)
+	}
+
+	if !reflect.DeepEqual(vmObj.Spec.Template.Spec.Domain.Devices, vmObjCopy.Spec.Template.Spec.Domain.Devices) {
+		logrus.WithFields(logrus.Fields{
+			"name":      vmi.Name,
+			"namespace": vmi.Namespace,
+		}).Debugf("updating vm device allocation details: %v", vmObj.Spec.Template.Spec.Domain.Devices)
+		_, err := h.vmClient.Update(vmObj)
+		return vmi, err
+	}
+	return vmi, nil
+}
+
+func patchHostDevices(vm *kubevirtv1.VirtualMachine, deviceInfo map[string][]string) {
+	for i, device := range vm.Spec.Template.Spec.Domain.Devices.HostDevices {
+		actualDevices, ok := deviceInfo[device.DeviceName]
+		if !ok {
+			continue // should not ideally be possible but in case it does happen we ignore and continue
+		}
+		// pop first element of the actualDevices and set to deviceNames
+		vm.Spec.Template.Spec.Domain.Devices.HostDevices[i].Name = actualDevices[0]
+		actualDevices = actualDevices[1:]
+		deviceInfo[device.DeviceName] = actualDevices // update map to ensure same name is not reused
+	}
+}
+
+func patchGPUDevices(vm *kubevirtv1.VirtualMachine, deviceInfo map[string][]string) {
+	for i, device := range vm.Spec.Template.Spec.Domain.Devices.GPUs {
+		actualDevices, ok := deviceInfo[device.DeviceName]
+		if !ok {
+			continue // should not ideally be possible but in case it does happen we ignore and continue
+		}
+		// pop first element of the actualDevices and set to deviceNames
+		vm.Spec.Template.Spec.Domain.Devices.GPUs[i].Name = actualDevices[0]
+		actualDevices = actualDevices[1:]
+		deviceInfo[device.DeviceName] = actualDevices // update map to ensure same name is not reused
+	}
+}

pkg/controller/virtualmachine/virtualmachine_test.go

@@ -0,0 +1,81 @@
+package virtualmachine
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	kubevirtv1 "kubevirt.io/api/core/v1"
+)
+
+func Test_patchHostDevices(t *testing.T) {
+	assert := require.New(t)
+
+	HostDevices := map[string][]string{
+		"device.com/sample": {"node1dev1", "node1dev2"},
+	}
+	vm := &kubevirtv1.VirtualMachine{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "demo",
+			Namespace: "default",
+		},
+		Spec: kubevirtv1.VirtualMachineSpec{
+			Template: &kubevirtv1.VirtualMachineInstanceTemplateSpec{
+				Spec: kubevirtv1.VirtualMachineInstanceSpec{
+					Domain: kubevirtv1.DomainSpec{
+						Devices: kubevirtv1.Devices{
+							HostDevices: []kubevirtv1.HostDevice{
+								{
+									Name:       "node1dev1",
+									DeviceName: "device.com/sample",
+								},
+								{
+									Name:       "randomDevice",
+									DeviceName: "device.com/sample",
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	patchHostDevices(vm, HostDevices)
+	// expect randomDevice to be replaced with actual device name
+	assert.Equal(vm.Spec.Template.Spec.Domain.Devices.HostDevices[1].Name, "node1dev2")
+
+}
+
+func Test_patchGPUDevices(t *testing.T) {
+	assert := require.New(t)
+
+	GPUDevices := map[string][]string{
+		"nvidia.com/A2-Q2": {"node1dev1"},
+	}
+	vm := &kubevirtv1.VirtualMachine{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      "demo",
+			Namespace: "default",
+		},
+		Spec: kubevirtv1.VirtualMachineSpec{
+			Template: &kubevirtv1.VirtualMachineInstanceTemplateSpec{
+				Spec: kubevirtv1.VirtualMachineInstanceSpec{
+					Domain: kubevirtv1.DomainSpec{
+						Devices: kubevirtv1.Devices{
+							GPUs: []kubevirtv1.GPU{
+								{
+									Name:       "sample",
+									DeviceName: "nvidia.com/A2-Q2",
+								},
+							},
+						},
+					},
+				},
+			},
+		},
+	}
+	patchGPUDevices(vm, GPUDevices)
+	// expect randomDevice to be replaced with actual device name
+	assert.Equal(vm.Spec.Template.Spec.Domain.Devices.GPUs[0].Name, "node1dev1")
+
+}