Merge pull request Project-HAMi#26 from DSFans2014/feat/set_numa_info

feat: set `NetworkID`

Author: archlitchi

cmd/main.go

@@ -23,6 +23,7 @@ import (
 	"os"
 	"syscall"
 
+	"github.com/Project-HAMi/HAMi/pkg/util/client"
 	"github.com/Project-HAMi/ascend-device-plugin/internal"
 	"github.com/Project-HAMi/ascend-device-plugin/internal/manager"
 	"github.com/Project-HAMi/ascend-device-plugin/internal/server"
@@ -136,6 +137,7 @@ func main() {
 	if err != nil {
 		klog.Fatalf("init PluginServer failed, error is %v", err)
 	}
+	client.InitGlobalClient()
 
 	err = start(server)
 	if err != nil {

config.yaml

@@ -57,3 +57,23 @@ vnpus:
       memory: 12288
       aiCore: 4
       aiCPU: 4
+- chipName: 910ProB
+  commonWord: Ascend910ProB
+  resourceName: huawei.com/Ascend910ProB
+  resourceMemoryName: huawei.com/Ascend910ProB-memory
+  memoryAllocatable: 32768
+  memoryCapacity: 32768
+  aiCore: 30
+  templates:
+    - name: vir02
+      memory: 2184
+      aiCore: 2
+    - name: vir04
+      memory: 4369
+      aiCore: 4
+    - name: vir08
+      memory: 8738
+      aiCore: 8
+    - name: vir16
+      memory: 17476
+      aiCore: 16

go.mod

@@ -57,6 +57,6 @@ require (
 )
 
 replace (
-	github.com/Project-HAMi/HAMi v0.0.0 => github.com/Project-HAMi/HAMi v0.0.0-20250107033239-d04fc8baaad6
+	github.com/Project-HAMi/HAMi v0.0.0 => github.com/Project-HAMi/HAMi v0.0.0-20250901013025-61c6cbe7d480
 	huawei.com/npu-exporter/v6 => gitee.com/ascend/ascend-npu-exporter/v6 v6.0.0-RC3
 )

internal/server/server.go

@@ -23,8 +23,10 @@ import (
 	"net"
 	"os"
 	"path"
+	"strings"
 	"time"
 
+	"github.com/Project-HAMi/HAMi/pkg/device"
 	"github.com/Project-HAMi/HAMi/pkg/device/ascend"
 	"github.com/Project-HAMi/HAMi/pkg/util"
 	"github.com/Project-HAMi/HAMi/pkg/util/nodelock"
@@ -40,7 +42,8 @@ import (
 const (
 	// RegisterAnnos = "hami.io/node-register-ascend"
 	// PodAllocAnno = "huawei.com/AscendDevices"
-	NodeLockAscend = "hami.io/mutex.lock"
+	NodeLockAscend  = "hami.io/mutex.lock"
+	Ascend910Prefix = "Ascend910"
 )
 
 var (
@@ -188,12 +191,19 @@ func (ps *PluginServer) registerKubelet() error {
 	return nil
 }
 
+func (ps *PluginServer) getDeviceNetworkID(idx int) (int, error) {
+	if idx > 3 {
+		return 1, nil
+	}
+	return 0, nil
+}
+
 func (ps *PluginServer) registerHAMi() error {
 	devs := ps.mgr.GetDevices()
-	apiDevices := make([]*util.DeviceInfo, 0, len(devs))
+	apiDevices := make([]*device.DeviceInfo, 0, len(devs))
 	// hami currently believes that the index starts from 0 and is continuous.
 	for i, dev := range devs {
-		apiDevices = append(apiDevices, &util.DeviceInfo{
+		device := &device.DeviceInfo{
 			Index:   uint(i),
 			ID:      dev.UUID,
 			Count:   int32(ps.mgr.VDeviceCount()),
@@ -202,10 +212,20 @@ func (ps *PluginServer) registerHAMi() error {
 			Type:    ps.mgr.CommonWord(),
 			Numa:    0,
 			Health:  dev.Health,
-		})
+		}
+		if strings.HasPrefix(device.Type, Ascend910Prefix) {
+			NetworkID, err := ps.getDeviceNetworkID(i)
+			if err != nil {
+				return fmt.Errorf("get networkID error: %v", err)
+			}
+			device.CustomInfo = map[string]any{
+				"NetworkID": NetworkID,
+			}
+		}
+		apiDevices = append(apiDevices, device)
 	}
 	annos := make(map[string]string)
-	annos[ps.registerAnno] = util.MarshalNodeDevices(apiDevices)
+	annos[ps.registerAnno] = device.MarshalNodeDevices(apiDevices)
 	annos[ps.handshakeAnno] = "Reported_" + time.Now().Add(time.Duration(*reportTimeOffset)*time.Second).Format("2006.01.02 15:04:05")
 	node, err := util.GetNode(ps.nodeName)
 	if err != nil {