Skip to content

feat(scheduler): add node nouse gpuuuid function #1206

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZhengW22 wants to merge 3 commits into
base: master
Choose a base branch
from
Open

feat(scheduler): add node nouse gpuuuid function #1206

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
8b6b85a
Commit disable gpu uuid.
ZhengW22 Feb 2, 2026
ac5aa79
feat: add dynamic handling for Ascend device annotations in node sche…
ZhengW22 Feb 11, 2026
7023fb3
Merge branch 'master' into master
ZhengW22 Apr 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 61 additions & 0 deletions pkg/scheduler/nodes.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,31 @@ package scheduler

import (
"fmt"
"strings"
"sync"

corev1 "k8s.io/api/core/v1"
"k8s.io/klog/v2"

"github.com/Project-HAMi/HAMi/pkg/device"
"github.com/Project-HAMi/HAMi/pkg/device/cambricon"
"github.com/Project-HAMi/HAMi/pkg/device/hygon"
"github.com/Project-HAMi/HAMi/pkg/device/kunlun"
"github.com/Project-HAMi/HAMi/pkg/device/metax"
"github.com/Project-HAMi/HAMi/pkg/device/mthreads"
"github.com/Project-HAMi/HAMi/pkg/device/nvidia"
"github.com/Project-HAMi/HAMi/pkg/scheduler/policy"
)

var vendorNoUseAnnoKeyMap = map[string][]string{
nvidia.GPUNoUseUUID: {nvidia.NvidiaGPUDevice},
cambricon.MLUNoUseUUID: {cambricon.CambriconMLUDevice},
hygon.DCUNoUseUUID: {hygon.HygonDCUDevice},
mthreads.MthreadsNoUseUUID: {mthreads.MthreadsGPUDevice},
metax.MetaxNoUseUUID: {metax.MetaxGPUDevice, metax.MetaxSGPUDevice},
Comment thread
ZhengW22 marked this conversation as resolved.

@gemini-code-assist gemini-code-assist Bot Feb 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The constants metax.MetaxNoUseUUID and metax.MetaxSGPUDevice are used here, but they are not defined in the pkg/device/metax/device.go file provided in the context. This will lead to a compilation error. Please ensure these constants are defined in the metax package. For consistency with other device packages, MetaxNoUseUUID could be defined as "metax-tech.com/nouse-gpuuuid".

@ZhengW22 ZhengW22 Feb 2, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is a gemini problem. You can find those constants in pkg/device/metax/protocol.go and pkg/device/metax/sdevice.go easily. Because metax contributors don't put all constants in one files.

kunlun.KunlunNoUseUUID: {kunlun.KunlunGPUDevice},
}
Comment on lines +38 to +45

@gemini-code-assist gemini-code-assist Bot Jan 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

While this map works, it introduces a tight coupling between the scheduler package and all specific device packages (nvidia, cambricon, etc.). Every time a new device type with a 'no-use' annotation is added, this central map will need to be modified.

A more decoupled and maintainable approach would be to have each device implementation provide this information. For example, you could add a method to the device.Devices interface:

// In pkg/device/devices.go
type Devices interface {
    // ...
    GetNoUseAnnotation() (key string, vendors []string)
}

The scheduler could then build this map dynamically at startup by iterating over the registered device implementations. This would make the system more extensible.


type NodeUsage struct {
Node *corev1.Node
Devices policy.DeviceUsageList
Expand Down Expand Up @@ -60,6 +76,51 @@ func (m *nodeManager) addNode(nodeID string, nodeInfo *device.NodeInfo) {
} else {
m.nodes[nodeID] = nodeInfo
}
m.nodes[nodeID].Devices = rmDeviceByNodeAnnotation(m.nodes[nodeID])
}

func rmDeviceByNodeAnnotation(nodeInfo *device.NodeInfo) map[string][]device.DeviceInfo {
if nodeInfo == nil {
return nil
}
vendorWithDisableGPUUUIDMap := make(map[string]map[string]bool)

@gemini-code-assist gemini-code-assist Bot Jan 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

The function rmDeviceByNodeAnnotation does not check if nodeInfo is nil. If a nil nodeInfo is passed, it could cause a panic at nodeInfo.Node. Although the current call site in addNode seems to prevent this, adding a nil check at the beginning of the function is a good practice for robustness and to make the function safer for future use.

	if nodeInfo == nil {
		return nil
	}
	vendorWithDisableGPUUUIDMap := make(map[string]map[string]bool)

if nodeInfo.Node != nil && nodeInfo.Node.Annotations != nil {
for annoKey, vendors := range vendorNoUseAnnoKeyMap {
klog.V(5).Infof("Current annokey is %s, and vendor is %v", annoKey, vendors)
if value, ok := nodeInfo.Node.Annotations[annoKey]; ok {
disableGPUUUIDList := strings.Split(value, ",")
klog.V(5).Infof("Disable gpu uuid list is: %v", disableGPUUUIDList)
for _, disableGPUUUID := range disableGPUUUIDList {
if id := strings.TrimSpace(disableGPUUUID); id != "" {
for _, vendor := range vendors {
if vendorWithDisableGPUUUIDMap[vendor] == nil {
vendorWithDisableGPUUUIDMap[vendor] = make(map[string]bool)
}
Comment thread
ZhengW22 marked this conversation as resolved.
vendorWithDisableGPUUUIDMap[vendor][id] = true
}
}
}
Comment thread
ZhengW22 marked this conversation as resolved.
}
}
}
if len(vendorWithDisableGPUUUIDMap) == 0 {
return nodeInfo.Devices
}
newDeviceMap := make(map[string][]device.DeviceInfo)
for deviceName, deviceList := range nodeInfo.Devices {
newDeviceList := make([]device.DeviceInfo, 0, len(deviceList))
for _, d := range deviceList {
if disableGPUUUIDMap, ok := vendorWithDisableGPUUUIDMap[d.DeviceVendor]; ok {
if disabled := disableGPUUUIDMap[d.ID]; disabled {
klog.V(5).Infof("Disable gpu uuid is : %s", d.ID)
continue
}
Comment on lines +137 to +140

@gemini-code-assist gemini-code-assist Bot Feb 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This conditional check can be simplified. In Go, when checking a boolean value from a map, you can directly use the result of the map access in an if statement. If the key doesn't exist, the expression will evaluate to false, which is the desired behavior here.

Suggested change
if disabled := disableGPUUUIDMap[d.ID]; disabled {
klog.V(5).Infof("Disable gpu uuid is : %s", d.ID)
continue
}
if disableGPUUUIDMap[d.ID] {
klog.V(5).Infof("Disable gpu uuid is : %s", d.ID)
continue
}

}
newDeviceList = append(newDeviceList, d)
}
newDeviceMap[deviceName] = newDeviceList
}
Comment on lines +132 to +145

@gemini-code-assist gemini-code-assist Bot Jan 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The logic for filtering devices is correct. However, for better performance and readability, the check for whether a vendor has devices to disable can be hoisted out of the inner loop. Since d.DeviceVendor will be the same as deviceName for all devices in deviceList, you can check for vendorWithDisableGPUUUIDMap[deviceName] once before iterating through the deviceList.

 newDeviceMap := make(map[string][]device.DeviceInfo, len(nodeInfo.Devices))
 for deviceName, deviceList := range nodeInfo.Devices {
  disableGPUUUIDMap, ok := vendorWithDisableGPUUUIDMap[deviceName]
  if !ok {
   newDeviceMap[deviceName] = deviceList
   continue
  }
  newDeviceList := make([]device.DeviceInfo, 0, len(deviceList))
  for _, d := range deviceList {
   if disabled := disableGPUUUIDMap[d.ID]; disabled {
    klog.V(5).Infof("Disable gpu uuid is : %s", d.ID)
    continue
   }
   newDeviceList = append(newDeviceList, d)
  }
  newDeviceMap[deviceName] = newDeviceList
 }

return newDeviceMap
}

func (m *nodeManager) rmNodeDevices(nodeID string, deviceVendor string) {
Expand Down
107 changes: 107 additions & 0 deletions pkg/scheduler/nodes_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,13 +18,18 @@ package scheduler

import (
"fmt"
"reflect"
"strings"
"testing"

"github.com/Project-HAMi/HAMi/pkg/device"
"github.com/Project-HAMi/HAMi/pkg/device/metax"
"github.com/Project-HAMi/HAMi/pkg/device/nvidia"
"github.com/Project-HAMi/HAMi/pkg/scheduler/config"

"gotest.tools/v3/assert"
corev1 "k8s.io/api/core/v1"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
)

func Test_addNode_ListNodes(t *testing.T) {
Expand Down Expand Up @@ -322,3 +327,105 @@ func Test_rmNodeDevices(t *testing.T) {
})
}
}

func Test_rmDeviceByNodeAnnotation(t *testing.T) {

@gemini-code-assist gemini-code-assist Bot Dec 8, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The test suite for rmDeviceByNodeAnnotation is good, but it would be beneficial to add a test case that verifies that UUIDs with leading/trailing whitespace in the annotation value are handled correctly. This would ensure the strings.TrimSpace logic is effective and prevent regressions.

For example:

{
    name: "Test remove device with whitespace in annotation",
    args: args{
        nodeInfo: &device.NodeInfo{
            Node:    &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: " " + id1 + " "}}},
            Devices: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}},
        },
    },
    want: []device.DeviceInfo{},
},

@gemini-code-assist gemini-code-assist Bot Jan 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The current tests don't cover cases where the comma-separated UUID list in the annotation contains spaces (e.g., "uuid1, uuid2"). This could hide a bug where strings.TrimSpace is not being used correctly when populating the map of disabled UUIDs. Please consider adding a test case to cover this scenario to make the tests more robust.

id1 := "60151478-4709-4242-a8c1-a944252d194b"
id2 := "33c00a52-72ab-4b61-a7ce-43107588835b"
type args struct {
nodeInfo *device.NodeInfo
}
tests := []struct {
name string
args args
want map[string][]device.DeviceInfo
}{
{
name: "Test space condition",
args: args{
nodeInfo: &device.NodeInfo{
Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: strings.Join([]string{id1 + " ", " " + id2 + " "}, ",")}}},
Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}, {DeviceVendor: nvidia.NvidiaGPUDevice, ID: id2}}},
},
},
want: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{}},
},
{
name: "Test remove one device",
args: args{
nodeInfo: &device.NodeInfo{
Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: id1}}},
Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}}},
},
},
want: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{}},
},
{
name: "Test remove two devices",
args: args{
nodeInfo: &device.NodeInfo{
Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: strings.Join([]string{id1, id2}, ",")}}},
Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}, {DeviceVendor: nvidia.NvidiaGPUDevice, ID: id2}}},
},
},
want: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{}},
},
{
name: "Test remove one device and keep one device",
args: args{
nodeInfo: &device.NodeInfo{
Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: strings.Join([]string{id2}, ",")}}},
Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}, {DeviceVendor: nvidia.NvidiaGPUDevice, ID: id2}}},
},
},
want: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}}},
},
{
name: "Test no removing device, case1",
args: args{
nodeInfo: &device.NodeInfo{
Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{"test-key": ""}}},
Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}}},
},
},
want: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}}},
},
{
name: "Test no removing device, case2",
args: args{
nodeInfo: &device.NodeInfo{
Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: id2}}},
Devices: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}}},
},
},
want: map[string][]device.DeviceInfo{nvidia.NvidiaGPUDevice: []device.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}}},
},
{
name: "Test removing metax device, case1",
args: args{
nodeInfo: &device.NodeInfo{
Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{metax.MetaxNoUseUUID: id1}}},
// Devices: []device.DeviceInfo{{DeviceVendor: metax.MetaxGPUDevice, ID: id1}},

@gemini-code-assist gemini-code-assist Bot Jan 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This commented-out line appears to be a remnant from development and can be safely removed to improve code clarity.

@gemini-code-assist gemini-code-assist Bot Feb 2, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

This commented-out line appears to be a remnant from development. It should be removed to improve code clarity and maintainability.

Devices: map[string][]device.DeviceInfo{metax.MetaxGPUDevice: []device.DeviceInfo{{DeviceVendor: metax.MetaxGPUDevice, ID: id1}}},
},
},
want: map[string][]device.DeviceInfo{metax.MetaxGPUDevice: []device.DeviceInfo{}},
},
{
name: "Test removing metax sgpu device",
args: args{
nodeInfo: &device.NodeInfo{
Node: &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{metax.MetaxNoUseUUID: id1}}},
Devices: map[string][]device.DeviceInfo{metax.MetaxSGPUDevice: {{DeviceVendor: metax.MetaxSGPUDevice, ID: id1}}},
},
},
want: map[string][]device.DeviceInfo{metax.MetaxSGPUDevice: []device.DeviceInfo{}},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if got := rmDeviceByNodeAnnotation(tt.args.nodeInfo); !reflect.DeepEqual(got, tt.want) {
t.Errorf("rmDeviceByNodeAnnotation() = %v, want %v", got, tt.want)
}
Comment on lines +456 to +458

@gemini-code-assist gemini-code-assist Bot Jan 30, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For consistency with other tests in this file (e.g., Test_addNode_ListNodes), it's better to use assert.DeepEqual for comparing the expected and actual results. This improves the maintainability and uniformity of the test suite.

   got := rmDeviceByNodeAnnotation(tt.args.nodeInfo)
   assert.DeepEqual(t, tt.want, got)

})
}
}
Comment on lines +331 to +461

@gemini-code-assist gemini-code-assist Bot Jul 31, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The test cases for rmDeviceByNodeAnnotation are not comprehensive enough. Please consider adding more tests to cover the following scenarios:

  • Multiple UUIDs in the annotation value.
  • Multiple devices on the node, with some being removed and some not.
  • Multiple annotations for different vendors on the same node.
  • Malformed annotation values (e.g., with extra spaces or empty parts like "uuid1,,uuid2").
  • A case where a UUID matches but the device vendor does not.