From 4888f81ff6a29ed6ca8b26e12c6399a91025176c Mon Sep 17 00:00:00 2001
From: wanyaoqi <d3lx.yq@gmail.com>
Date: Mon, 11 Sep 2023 14:50:07 +0800
Subject: [PATCH] fix(region): check nvidia vgpu count and conflict with other
 gpus

Signed-off-by: wanyaoqi <d3lx.yq@gmail.com>
---
 pkg/compute/models/guest_actions.go | 29 +++++++++++++++++++++++++++++
 pkg/compute/models/guests.go        | 17 +++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/pkg/compute/models/guest_actions.go b/pkg/compute/models/guest_actions.go
index 2a7424b5ec5..e2131ca43b3 100644
--- a/pkg/compute/models/guest_actions.go
+++ b/pkg/compute/models/guest_actions.go
@@ -2046,6 +2046,20 @@ func (self *SGuest) startAttachIsolatedDevices(ctx context.Context, userCred mcc
 		}
 	}
 
+	if dev.DevType == api.LEGACY_VGPU_TYPE {
+		devs, err := self.GetIsolatedDevices()
+		if err != nil {
+			return errors.Wrap(err, "get isolated devices")
+		}
+		for i := range devs {
+			if devs[i].DevType == api.LEGACY_VGPU_TYPE {
+				return httperrors.NewBadRequestError("Nvidia vgpu count exceed > 1")
+			} else if utils.IsInStringArray(devs[i].DevType, api.VALID_GPU_TYPES) {
+				return httperrors.NewBadRequestError("Nvidia vgpu can't passthrough with other gpus")
+			}
+		}
+	}
+
 	defer func() { go host.ClearSchedDescCache() }()
 	for i := 0; i < len(devs); i++ {
 		err = self.attachIsolatedDevice(ctx, userCred, &devs[i], nil, nil)
@@ -2085,6 +2099,21 @@ func (self *SGuest) startAttachIsolatedDevGeneral(ctx context.Context, userCred
 	if !utils.IsInStringArray(self.GetStatus(), []string{api.VM_READY, api.VM_RUNNING}) {
 		return httperrors.NewInvalidStatusError("Can't attach GPU when status is %q", self.GetStatus())
 	}
+
+	if dev.DevType == api.LEGACY_VGPU_TYPE {
+		devs, err := self.GetIsolatedDevices()
+		if err != nil {
+			return errors.Wrap(err, "get isolated devices")
+		}
+		for i := range devs {
+			if devs[i].DevType == api.LEGACY_VGPU_TYPE {
+				return httperrors.NewBadRequestError("Nvidia vgpu count exceed > 1")
+			} else if utils.IsInStringArray(devs[i].DevType, api.VALID_GPU_TYPES) {
+				return httperrors.NewBadRequestError("Nvidia vgpu can't passthrough with other gpus")
+			}
+		}
+	}
+
 	host, _ := self.GetHost()
 	lockman.LockObject(ctx, host)
 	defer lockman.ReleaseObject(ctx, host)
diff --git a/pkg/compute/models/guests.go b/pkg/compute/models/guests.go
index 5103fa34d21..46998aef388 100644
--- a/pkg/compute/models/guests.go
+++ b/pkg/compute/models/guests.go
@@ -1828,6 +1828,23 @@ func (manager *SGuestManager) validateCreateData(
 		input.IsolatedDevices[idx] = devConfig
 	}
 
+	nvidiaVgpuCnt := 0
+	gpuCnt := 0
+	for i := 0; i < len(input.IsolatedDevices); i++ {
+		if input.IsolatedDevices[i].DevType == api.LEGACY_VGPU_TYPE {
+			nvidiaVgpuCnt += 1
+		} else if utils.IsInStringArray(input.IsolatedDevices[i].DevType, api.VALID_GPU_TYPES) {
+			gpuCnt += 1
+		}
+	}
+
+	if nvidiaVgpuCnt > 1 {
+		return nil, httperrors.NewBadRequestError("Nvidia vgpu count exceed > 1")
+	}
+	if nvidiaVgpuCnt > 0 && gpuCnt > 0 {
+		return nil, httperrors.NewBadRequestError("Nvidia vgpu can't passthrough with other gpus")
+	}
+
 	keypairId := input.KeypairId
 	if len(keypairId) > 0 {
 		keypairObj, err := KeypairManager.FetchByIdOrName(userCred, keypairId)