From 863432421f56ad7116e6be670bd72c72aaff9edf Mon Sep 17 00:00:00 2001
From: Nimbus318 <136771156+Nimbus318@users.noreply.github.com>
Date: Tue, 14 Jan 2025 13:17:51 +0800
Subject: [PATCH] fix: associate resource allocation between MIG-enabled cards
 and tasks

Signed-off-by: Nimbus318 <136771156+Nimbus318@users.noreply.github.com>
---
 server/internal/biz/pod.go           | 3 ++-
 server/internal/exporter/exporter.go | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/server/internal/biz/pod.go b/server/internal/biz/pod.go
index 8751071..15aa6c5 100644
--- a/server/internal/biz/pod.go
+++ b/server/internal/biz/pod.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"github.com/go-kratos/kratos/v2/log"
 	k8stypes "k8s.io/apimachinery/pkg/types"
+	"strings"
 	"time"
 )
 
@@ -76,7 +77,7 @@ func ContainersStatisticsInfo(containers []*Container, deviceId string) (int32,
 	var memory int32 = 0
 	for _, t := range containers {
 		for _, cd := range t.ContainerDevices {
-			if deviceId != "" && cd.UUID != deviceId {
+			if deviceId != "" && !strings.HasPrefix(cd.UUID, deviceId) {
 				continue
 			}
 			vGPU = vGPU + 1
diff --git a/server/internal/exporter/exporter.go b/server/internal/exporter/exporter.go
index bd27668..cd341c8 100644
--- a/server/internal/exporter/exporter.go
+++ b/server/internal/exporter/exporter.go
@@ -5,6 +5,7 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"strings"
 	"time"
 	pb "vgpu/api/v1"
 	"vgpu/internal/biz"
@@ -163,7 +164,7 @@ func (s *MetricsGenerator) GenerateContainerMetrics(ctx context.Context) error {
 			var memory int32 = 0
 			var provider string = ""
 			for _, cd := range c.ContainerDevices {
-				if device.AliasId != "" && device.AliasId != cd.UUID {
+				if device.AliasId != "" && !strings.HasPrefix(cd.UUID, device.AliasId) {
 					continue
 				}
 				vGPU = vGPU + 1