Support Metax sGPU topology aware (#1193)

Signed-off-by: Lei Guo <[email protected]>
Co-authored-by: Lei Guo <[email protected]>

Author: Kyrie336

charts/hami/templates/scheduler/device-configmap.yaml

@@ -97,6 +97,7 @@ data:
       resourceVCountName: {{ .Values.metaxResourceName }}
       resourceVMemoryName: {{ .Values.metaxResourceMem }}
       resourceVCoreName: {{ .Values.metaxResourceCore }}
+      sgpuTopologyAware: {{ .Values.metaxsGPUTopologyAware }}
     enflame:
       resourceCountName: "enflame.com/vgcu"
       resourcePercentageName: "enflame.com/vgcu-percentage"

charts/hami/values.yaml

@@ -43,10 +43,11 @@ iluvatarResourceName: "iluvatar.ai/vgpu"
 iluvatarResourceMem: "iluvatar.ai/vcuda-memory"
 iluvatarResourceCore: "iluvatar.ai/vcuda-core"
 
-#Metax SGPU Parameters
+#Metax sGPU Parameters
 metaxResourceName: "metax-tech.com/sgpu"
 metaxResourceCore: "metax-tech.com/vcore"
 metaxResourceMem: "metax-tech.com/vmemory"
+metaxsGPUTopologyAware: "false"
 
 #Kunlun XPU Parameters
 kunlunResourceName: "kunlunxin.com/xpu"

examples/metax/sgpu/allocate_specific_gpu.yaml

@@ -3,7 +3,7 @@ kind: Pod
 metadata:
   name: gpu-pod
   annotations:
-    metax-tech.com/use-gpuuuid: "36beae85-c835-6b14-6ab2-02671837a59c" # allocate specific gpu
+    metax-tech.com/use-gpuuuid: "GPU-36beae85-c835-6b14-6ab2-02671837a59c" # allocate specific gpu
 spec:
   containers:
     - name: ubuntu-container

examples/metax/sgpu/enable_topology_aware.yaml

@@ -0,0 +1,15 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: gpu-pod
+  annotations:
+    metax-tech.com/sgpu-topology-aware: "true" # enable topology aware scheduling
+spec:
+  containers:
+    - name: ubuntu-container
+      image: ubuntu:22.04
+      imagePullPolicy: IfNotPresent
+      command: ["sleep","infinity"]
+      resources:
+        limits:
+          metax-tech.com/sgpu: 4 # requesting 4 exclusive GPU

pkg/device/metax/config.go

@@ -26,6 +26,7 @@ type MetaxConfig struct {
 	ResourceVCountName  string `yaml:"resourceVCountName"`
 	ResourceVMemoryName string `yaml:"resourceVMemoryName"`
 	ResourceVCoreName   string `yaml:"resourceVCoreName"`
+	TopologyAware       bool   `yaml:"sgpuTopologyAware"`
 }
 
 func ParseConfig(fs *flag.FlagSet) {
@@ -36,4 +37,5 @@ func ParseConfig(fs *flag.FlagSet) {
 	fs.StringVar(&MetaxResourceNameVCount, "metax-vcount", "metax-tech.com/sgpu", "metax vcount name")
 	fs.StringVar(&MetaxResourceNameVCore, "metax-vcore", "metax-tech.com/vcore", "metax vcore name")
 	fs.StringVar(&MetaxResourceNameVMemory, "metax-vmemory", "metax-tech.com/vmemory", "metax vmemory name")
+	fs.BoolVar(&MetaxTopologyAware, "sgpu-topology-aware", false, "sGPU topology aware enable")
 }

pkg/device/metax/protocol.go

@@ -30,7 +30,8 @@ const (
 	MetaxUseUUID   = "metax-tech.com/use-gpuuuid"
 	MetaxNoUseUUID = "metax-tech.com/nouse-gpuuuid"
 
-	MetaxSGPUQosPolicy = "metax-tech.com/sgpu-qos-policy"
+	MetaxSGPUQosPolicy     = "metax-tech.com/sgpu-qos-policy"
+	MetaxSGPUTopologyAware = "metax-tech.com/sgpu-topology-aware"
 )
 
 const (
@@ -52,6 +53,7 @@ type MetaxSDeviceInfo struct {
 	Numa              int32  `json:"numa,omitempty"`
 	Healthy           bool   `json:"healthy,omitempty"`
 	QosPolicy         string `json:"qosPolicy,omitempty"`
+	LinkZone          int32  `json:"linkZone,omitempty"`
 }
 type NodeMetaxSDeviceInfo []*MetaxSDeviceInfo
 
@@ -67,8 +69,8 @@ func (ni NodeMetaxSDeviceInfo) String() string {
 	str := "\n"
 
 	for _, i := range ni {
-		str += fmt.Sprintf("MetaxSDeviceInfo[%s]: TotalDevCount=%d, TotalCompute=%d, TotalVRam=%d, Numa=%d, Healthy=%t, QosPolicy=%s\n",
-			i.UUID, i.TotalDevCount, i.TotalCompute, i.TotalVRam, i.Numa, i.Healthy, i.QosPolicy)
+		str += fmt.Sprintf("MetaxSDeviceInfo[%s]: TotalDevCount=%d, TotalCompute=%d, TotalVRam=%d, Numa=%d, Healthy=%t, QosPolicy=%s, LinkZone=%d\n",
+			i.UUID, i.TotalDevCount, i.TotalCompute, i.TotalVRam, i.Numa, i.Healthy, i.QosPolicy, i.LinkZone)
 	}
 
 	return str
@@ -108,6 +110,7 @@ func convertMetaxSDeviceToHAMIDevice(metaxSDevices []*MetaxSDeviceInfo) []*util.
 			CustomInfo: map[string]any{
 				"QosPolicy": sdevice.QosPolicy,
 				"Model":     sdevice.Model,
+				"LinkZone":  sdevice.LinkZone,
 			},
 		}
 	}

pkg/device/metax/protocol_test.go

@@ -46,6 +46,7 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 					Numa:              0,
 					Healthy:           true,
 					QosPolicy:         BestEffort,
+					LinkZone:          1,
 				},
 				{
 					UUID:              "GPU-a16ac188-0592-5c8f-2b6e-8bd8e7a604a8",
@@ -60,6 +61,7 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 					Numa:              -1,
 					Healthy:           false,
 					QosPolicy:         BurstShare,
+					LinkZone:          2,
 				},
 			},
 			expected: []*util.DeviceInfo{
@@ -78,6 +80,7 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 					CustomInfo: map[string]any{
 						"QosPolicy": BestEffort,
 						"Model":     "native",
+						"LinkZone":  int32(1),
 					},
 				},
 				{
@@ -95,6 +98,7 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 					CustomInfo: map[string]any{
 						"QosPolicy": BurstShare,
 						"Model":     "sgpu",
+						"LinkZone":  int32(2),
 					},
 				},
 			},

pkg/device/metax/score.go

@@ -0,0 +1,66 @@
+/*
+Copyright 2025 The HAMi Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package metax
+
+import "fmt"
+
+const DirectLinkScore = 10
+
+type LinkDevice struct {
+	uuid     string
+	linkZone int32
+}
+
+func (from *LinkDevice) score(to *LinkDevice) int {
+	if from.uuid == to.uuid {
+		return 0
+	}
+
+	if from.linkZone == 0 || to.linkZone == 0 {
+		return 0
+	}
+
+	if from.linkZone == to.linkZone {
+		return DirectLinkScore
+	} else {
+		return 0
+	}
+}
+
+type LinkDevices []*LinkDevice
+
+func (devs LinkDevices) Score() int {
+	score := 0
+
+	for i := 0; i < len(devs); i++ {
+		for j := i + 1; j < len(devs); j++ {
+			score += devs[i].score(devs[j])
+		}
+	}
+
+	return score
+}
+
+func (devs LinkDevices) String() string {
+	str := "["
+	for _, dev := range devs {
+		str += fmt.Sprintf("%v", *dev)
+	}
+	str += "]"
+
+	return str
+}