Support Metax sGPU Qos Policy (#1123)

Kyrie336 · Lei Guo · web-flow · commit 7f0dadc868f4 · 2025-06-20T15:54:52.000+08:00
Signed-off-by: Lei Guo &lt;Lei.Guo@metax-tech.com&gt;
Co-authored-by: Lei Guo &lt;Lei.Guo@metax-tech.com&gt;
diff --git a/docs/metax-support.md b/docs/metax-support.md
@@ -40,7 +40,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04 
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/docs/metax-support_cn.md b/docs/metax-support_cn.md
@@ -38,7 +38,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04 
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/examples/metax/gpu/binpack.yaml b/examples/metax/gpu/binpack.yaml
@@ -7,7 +7,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04 
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/examples/metax/gpu/default_use.yaml b/examples/metax/gpu/default_use.yaml
@@ -5,7 +5,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/examples/metax/gpu/spread.yaml b/examples/metax/gpu/spread.yaml
@@ -7,7 +7,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04 
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/examples/metax/sgpu/allocate_exclusive.yaml b/examples/metax/sgpu/allocate_exclusive.yaml
@@ -5,7 +5,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/examples/metax/sgpu/allocate_specific_gpu.yaml b/examples/metax/sgpu/allocate_specific_gpu.yaml
@@ -7,7 +7,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/examples/metax/sgpu/allocate_specific_qos.yaml b/examples/metax/sgpu/allocate_specific_qos.yaml
@@ -0,0 +1,17 @@
+apiVersion: v1
+kind: Pod
+metadata:
+  name: gpu-pod
+  annotations:
+    metax-tech.com/sgpu-qos-policy: "best-effort" # allocate specific qos sgpu
+spec:
+  containers:
+    - name: ubuntu-container
+      image: ubuntu:22.04
+      imagePullPolicy: IfNotPresent
+      command: ["sleep","infinity"]
+      resources:
+        limits:
+          metax-tech.com/sgpu: 1 # requesting 1 GPU
+          metax-tech.com/vcore: 60 # each GPU use 60% of total compute cores
+          metax-tech.com/vmemory: 4 # each GPU require 4 GiB device memory
diff --git a/examples/metax/sgpu/allocate_vmemory_MiB.yaml b/examples/metax/sgpu/allocate_vmemory_MiB.yaml
@@ -5,7 +5,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/examples/metax/sgpu/default_use.yaml b/examples/metax/sgpu/default_use.yaml
@@ -5,7 +5,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/examples/metax/sgpu/multi-containers.yaml b/examples/metax/sgpu/multi-containers.yaml
@@ -5,7 +5,7 @@ metadata:
 spec:
   containers:
     - name: ubuntu-container-1
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
@@ -14,7 +14,7 @@ spec:
           metax-tech.com/vcore: 60 # each GPU use 60% of total compute cores
           metax-tech.com/vmemory: 4 # each GPU require 4 GiB device memory
     - name: ubuntu-container-2
-      image: cr.metax-tech.com/public-ai-release/c500/colossalai:2.24.0.5-py38-ubuntu20.04-amd64 
+      image: ubuntu:22.04
       imagePullPolicy: IfNotPresent
       command: ["sleep","infinity"]
       resources:
diff --git a/pkg/device/ascend/device.go b/pkg/device/ascend/device.go
@@ -144,7 +144,7 @@ func (dev *Devices) GetNodeDevices(n corev1.Node) ([]*util.DeviceInfo, error) {
 	return nodeDevices, nil
 }
 
-func (dev *Devices) PatchAnnotations(annoInput *map[string]string, pd util.PodDevices) map[string]string {
+func (dev *Devices) PatchAnnotations(pod *corev1.Pod, annoInput *map[string]string, pd util.PodDevices) map[string]string {
 	commonWord := dev.CommonWord()
 	devList, ok := pd[commonWord]
 	if ok && len(devList) > 0 {
@@ -281,7 +281,7 @@ func (dev *Devices) ScoreNode(node *corev1.Node, podDevices util.PodSingleDevice
 	return 0
 }
 
-func (dev *Devices) AddResourceUsage(n *util.DeviceUsage, ctr *util.ContainerDevice) error {
+func (dev *Devices) AddResourceUsage(pod *corev1.Pod, n *util.DeviceUsage, ctr *util.ContainerDevice) error {
 	n.Used++
 	n.Usedcores += ctr.Usedcores
 	n.Usedmem += ctr.Usedmem
diff --git a/pkg/device/ascend/device_test.go b/pkg/device/ascend/device_test.go
@@ -332,7 +332,7 @@ func Test_PatchAnnotations(t *testing.T) {
 	}
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
-			result := dev.PatchAnnotations(&test.args.annoinput, test.args.pd)
+			result := dev.PatchAnnotations(&corev1.Pod{}, &test.args.annoinput, test.args.pd)
 
 			assert.Equal(t, len(test.want), len(result), "Expected length of result to match want")
 			for k, v := range test.want {
diff --git a/pkg/device/cambricon/device.go b/pkg/device/cambricon/device.go
@@ -299,7 +299,7 @@ func (dev *CambriconDevices) GenerateResourceRequests(ctr *corev1.Container) uti
 	}
 }
 
-func (dev *CambriconDevices) PatchAnnotations(annoinput *map[string]string, pd util.PodDevices) map[string]string {
+func (dev *CambriconDevices) PatchAnnotations(pod *corev1.Pod, annoinput *map[string]string, pd util.PodDevices) map[string]string {
 	devlist, ok := pd[CambriconMLUDevice]
 	if ok {
 		(*annoinput)[DsmluResourceAssigned] = "false"
@@ -318,7 +318,7 @@ func (dev *CambriconDevices) ScoreNode(node *corev1.Node, podDevices util.PodSin
 	return 0
 }
 
-func (dev *CambriconDevices) AddResourceUsage(n *util.DeviceUsage, ctr *util.ContainerDevice) error {
+func (dev *CambriconDevices) AddResourceUsage(pod *corev1.Pod, n *util.DeviceUsage, ctr *util.ContainerDevice) error {
 	n.Used++
 	n.Usedcores += ctr.Usedcores
 	n.Usedmem += ctr.Usedmem
diff --git a/pkg/device/cambricon/device_test.go b/pkg/device/cambricon/device_test.go
@@ -443,7 +443,7 @@ func Test_PatchAnnotations(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			dev := CambriconDevices{}
-			result := dev.PatchAnnotations(&test.args.annoinput, test.args.pd)
+			result := dev.PatchAnnotations(&corev1.Pod{}, &test.args.annoinput, test.args.pd)
 			assert.Equal(t, len(test.want), len(result), "Expected length of result to match want")
 			for k, v := range test.want {
 				assert.Equal(t, v, result[k], "pod add annotation key [%s], values is [%s]", k, result[k])
diff --git a/pkg/device/devices.go b/pkg/device/devices.go
@@ -51,9 +51,9 @@ type Devices interface {
 	LockNode(n *corev1.Node, p *corev1.Pod) error
 	ReleaseNodeLock(n *corev1.Node, p *corev1.Pod) error
 	GenerateResourceRequests(ctr *corev1.Container) util.ContainerDeviceRequest
-	PatchAnnotations(annoinput *map[string]string, pd util.PodDevices) map[string]string
+	PatchAnnotations(pod *corev1.Pod, annoinput *map[string]string, pd util.PodDevices) map[string]string
 	ScoreNode(node *corev1.Node, podDevices util.PodSingleDevice, previous []*util.DeviceUsage, policy string) float32
-	AddResourceUsage(n *util.DeviceUsage, ctr *util.ContainerDevice) error
+	AddResourceUsage(pod *corev1.Pod, n *util.DeviceUsage, ctr *util.ContainerDevice) error
 	Fit(devices []*util.DeviceUsage, request util.ContainerDeviceRequest, annos map[string]string, pod *corev1.Pod, allocated *util.PodDevices) (bool, map[string]util.ContainerDevices, string)
 	// This should not be associated with a specific device object
 	//ParseConfig(fs *flag.FlagSet)
diff --git a/pkg/device/enflame/device.go b/pkg/device/enflame/device.go
@@ -138,7 +138,7 @@ func (dev *EnflameDevices) GetNodeDevices(n corev1.Node) ([]*util.DeviceInfo, er
 	return nodedevices, nil
 }
 
-func (dev *EnflameDevices) PatchAnnotations(annoinput *map[string]string, pd util.PodDevices) map[string]string {
+func (dev *EnflameDevices) PatchAnnotations(pod *corev1.Pod, annoinput *map[string]string, pd util.PodDevices) map[string]string {
 	devlist, ok := pd[EnflameGPUDevice]
 	if ok && len(devlist) > 0 {
 		(*annoinput)[util.SupportDevices[EnflameGPUDevice]] = util.EncodePodSingleDevice(devlist)
@@ -243,7 +243,7 @@ func (dev *EnflameDevices) ScoreNode(node *corev1.Node, podDevices util.PodSingl
 	return 0
 }
 
-func (dev *EnflameDevices) AddResourceUsage(n *util.DeviceUsage, ctr *util.ContainerDevice) error {
+func (dev *EnflameDevices) AddResourceUsage(pod *corev1.Pod, n *util.DeviceUsage, ctr *util.ContainerDevice) error {
 	n.Used++
 	n.Usedcores += ctr.Usedcores
 	n.Usedmem += ctr.Usedmem
diff --git a/pkg/device/enflame/device_test.go b/pkg/device/enflame/device_test.go
@@ -145,7 +145,7 @@ func TestPatchAnnotations(t *testing.T) {
 			}
 
 			dev := &EnflameDevices{}
-			got := dev.PatchAnnotations(&annoInputCopy, tt.podDevices)
+			got := dev.PatchAnnotations(&corev1.Pod{}, &annoInputCopy, tt.podDevices)
 
 			if len(got) != len(tt.expected) {
 				t.Errorf("PatchAnnotations() got %d annotations, expected %d", len(got), len(tt.expected))
diff --git a/pkg/device/hygon/device.go b/pkg/device/hygon/device.go
@@ -253,7 +253,7 @@ func (dev *DCUDevices) GenerateResourceRequests(ctr *corev1.Container) util.Cont
 	return util.ContainerDeviceRequest{}
 }
 
-func (dev *DCUDevices) PatchAnnotations(annoinput *map[string]string, pd util.PodDevices) map[string]string {
+func (dev *DCUDevices) PatchAnnotations(pod *corev1.Pod, annoinput *map[string]string, pd util.PodDevices) map[string]string {
 	devlist, ok := pd[HygonDCUDevice]
 	if ok && len(devlist) > 0 {
 		deviceStr := util.EncodePodSingleDevice(devlist)
@@ -269,7 +269,7 @@ func (dev *DCUDevices) ScoreNode(node *corev1.Node, podDevices util.PodSingleDev
 	return 0
 }
 
-func (dev *DCUDevices) AddResourceUsage(n *util.DeviceUsage, ctr *util.ContainerDevice) error {
+func (dev *DCUDevices) AddResourceUsage(pod *corev1.Pod, n *util.DeviceUsage, ctr *util.ContainerDevice) error {
 	n.Used++
 	n.Usedcores += ctr.Usedcores
 	n.Usedmem += ctr.Usedmem
diff --git a/pkg/device/hygon/device_test.go b/pkg/device/hygon/device_test.go
@@ -572,7 +572,7 @@ func Test_PatchAnnotations(t *testing.T) {
 	for _, test := range tests {
 		t.Run(test.name, func(t *testing.T) {
 			dev := DCUDevices{}
-			result := dev.PatchAnnotations(test.args.annoinput, test.args.pd)
+			result := dev.PatchAnnotations(&corev1.Pod{}, test.args.annoinput, test.args.pd)
 			assert.DeepEqual(t, result, test.want)
 		})
 	}
diff --git a/pkg/device/iluvatar/device.go b/pkg/device/iluvatar/device.go
@@ -110,7 +110,7 @@ func (dev *IluvatarDevices) GetNodeDevices(n corev1.Node) ([]*util.DeviceInfo, e
 	return nodedevices, nil
 }
 
-func (dev *IluvatarDevices) PatchAnnotations(annoinput *map[string]string, pd util.PodDevices) map[string]string {
+func (dev *IluvatarDevices) PatchAnnotations(pod *corev1.Pod, annoinput *map[string]string, pd util.PodDevices) map[string]string {
 	devlist, ok := pd[IluvatarGPUDevice]
 	if ok && len(devlist) > 0 {
 		(*annoinput)[util.InRequestDevices[IluvatarGPUDevice]] = util.EncodePodSingleDevice(devlist)
@@ -229,7 +229,7 @@ func (dev *IluvatarDevices) ScoreNode(node *corev1.Node, podDevices util.PodSing
 	return 0
 }
 
-func (dev *IluvatarDevices) AddResourceUsage(n *util.DeviceUsage, ctr *util.ContainerDevice) error {
+func (dev *IluvatarDevices) AddResourceUsage(pod *corev1.Pod, n *util.DeviceUsage, ctr *util.ContainerDevice) error {
 	n.Used++
 	n.Usedcores += ctr.Usedcores
 	n.Usedmem += ctr.Usedmem
diff --git a/pkg/device/iluvatar/device_test.go b/pkg/device/iluvatar/device_test.go
@@ -147,7 +147,7 @@ func TestPatchAnnotations(t *testing.T) {
 			maps.Copy(annoInputCopy, tt.annoInput)
 
 			dev := &IluvatarDevices{}
-			got := dev.PatchAnnotations(&annoInputCopy, tt.podDevices)
+			got := dev.PatchAnnotations(&corev1.Pod{}, &annoInputCopy, tt.podDevices)
 
 			if len(got) != len(tt.expected) {
 				t.Errorf("PatchAnnotations() got %d annotations, expected %d", len(got), len(tt.expected))
diff --git a/pkg/device/metax/device.go b/pkg/device/metax/device.go
@@ -82,7 +82,7 @@ func (dev *MetaxDevices) GetNodeDevices(n corev1.Node) ([]*util.DeviceInfo, erro
 	return nodedevices, nil
 }
 
-func (dev *MetaxDevices) PatchAnnotations(annoinput *map[string]string, pd util.PodDevices) map[string]string {
+func (dev *MetaxDevices) PatchAnnotations(pod *corev1.Pod, annoinput *map[string]string, pd util.PodDevices) map[string]string {
 	devlist, ok := pd[MetaxGPUDevice]
 	if ok && len(devlist) > 0 {
 		deviceStr := util.EncodePodSingleDevice(devlist)
@@ -207,7 +207,7 @@ func (dev *MetaxDevices) ScoreNode(node *corev1.Node, podDevices util.PodSingleD
 	return res
 }
 
-func (dev *MetaxDevices) AddResourceUsage(n *util.DeviceUsage, ctr *util.ContainerDevice) error {
+func (dev *MetaxDevices) AddResourceUsage(pod *corev1.Pod, n *util.DeviceUsage, ctr *util.ContainerDevice) error {
 	n.Used++
 	n.Usedcores += ctr.Usedcores
 	n.Usedmem += ctr.Usedmem
diff --git a/pkg/device/metax/protocol.go b/pkg/device/metax/protocol.go
@@ -29,6 +29,14 @@ const (
 
 	MetaxUseUUID   = "metax-tech.com/use-gpuuuid"
 	MetaxNoUseUUID = "metax-tech.com/nouse-gpuuuid"
+
+	MetaxSGPUQosPolicy = "metax-tech.com/sgpu-qos-policy"
+)
+
+const (
+	BestEffort = "best-effort"
+	FixedShare = "fixed-share"
+	BurstShare = "burst-share"
 )
 
 type MetaxSDeviceInfo struct {
@@ -43,6 +51,7 @@ type MetaxSDeviceInfo struct {
 	AvailableVRam     int32  `json:"availableVRam,omitempty"`
 	Numa              int32  `json:"numa,omitempty"`
 	Healthy           bool   `json:"healthy,omitempty"`
+	QosPolicy         string `json:"qosPolicy,omitempty"`
 }
 type NodeMetaxSDeviceInfo []*MetaxSDeviceInfo
 
@@ -58,8 +67,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\n",
-			i.UUID, i.TotalDevCount, i.TotalCompute, i.TotalVRam, i.Numa, i.Healthy)
+		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)
 	}
 
 	return str
@@ -96,6 +105,10 @@ func convertMetaxSDeviceToHAMIDevice(metaxSDevices []*MetaxSDeviceInfo) []*util.
 			MIGTemplate:  []util.Geometry{},
 			Health:       sdevice.Healthy,
 			DeviceVendor: MetaxSGPUDevice,
+			CustomInfo: map[string]any{
+				"QosPolicy": sdevice.QosPolicy,
+				"Model":     sdevice.Model,
+			},
 		}
 	}
 
diff --git a/pkg/device/metax/protocol_test.go b/pkg/device/metax/protocol_test.go
@@ -36,7 +36,7 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 				{
 					UUID:              "GPU-a16ac188-0592-5c8f-2b6e-8bd8e7a604a9",
 					BDF:               "0000:0a:00.0",
-					Model:             "",
+					Model:             "native",
 					TotalDevCount:     16,
 					TotalCompute:      100,
 					TotalVRam:         64 * 1024,
@@ -45,11 +45,12 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 					AvailableVRam:     32 * 1024,
 					Numa:              0,
 					Healthy:           true,
+					QosPolicy:         BestEffort,
 				},
 				{
 					UUID:              "GPU-a16ac188-0592-5c8f-2b6e-8bd8e7a604a8",
 					BDF:               "0000:09:00.0",
-					Model:             "",
+					Model:             "sgpu",
 					TotalDevCount:     8,
 					TotalCompute:      100,
 					TotalVRam:         32 * 1024,
@@ -58,6 +59,7 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 					AvailableVRam:     32 * 1024,
 					Numa:              -1,
 					Healthy:           false,
+					QosPolicy:         BurstShare,
 				},
 			},
 			expected: []*util.DeviceInfo{
@@ -73,6 +75,10 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 					MIGTemplate:  []util.Geometry{},
 					Health:       true,
 					DeviceVendor: MetaxSGPUDevice,
+					CustomInfo: map[string]any{
+						"QosPolicy": BestEffort,
+						"Model":     "native",
+					},
 				},
 				{
 					ID:           "GPU-a16ac188-0592-5c8f-2b6e-8bd8e7a604a8",
@@ -86,6 +92,10 @@ func TestConvertMetaxSDeviceToHAMIDevice(t *testing.T) {
 					MIGTemplate:  []util.Geometry{},
 					Health:       false,
 					DeviceVendor: MetaxSGPUDevice,
+					CustomInfo: map[string]any{
+						"QosPolicy": BurstShare,
+						"Model":     "sgpu",
+					},
 				},
 			},
 		},
diff --git a/pkg/device/metax/qos.go b/pkg/device/metax/qos.go
diff --git a/pkg/device/metax/qos_test.go b/pkg/device/metax/qos_test.go
diff --git a/pkg/device/metax/sdevice.go b/pkg/device/metax/sdevice.go
diff --git a/pkg/device/metax/sdevice_test.go b/pkg/device/metax/sdevice_test.go
diff --git a/pkg/device/mthreads/device.go b/pkg/device/mthreads/device.go
diff --git a/pkg/device/mthreads/device_test.go b/pkg/device/mthreads/device_test.go
diff --git a/pkg/device/nvidia/device.go b/pkg/device/nvidia/device.go
diff --git a/pkg/device/nvidia/device_test.go b/pkg/device/nvidia/device_test.go
diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
diff --git a/pkg/scheduler/score.go b/pkg/scheduler/score.go
diff --git a/pkg/util/types.go b/pkg/util/types.go
diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go
