refactor(server): split server package into multiple files

Author: houyuxi

internal/server/allocate.go

@@ -0,0 +1,189 @@
+package server
+
+import (
+	"encoding/json"
+	"fmt"
+	"strconv"
+
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/klog/v2"
+	"k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
+
+	"github.com/Project-HAMi/HAMi/pkg/device"
+	"github.com/Project-HAMi/HAMi/pkg/device-plugin/nvidiadevice/nvinternal/plugin"
+	"github.com/Project-HAMi/HAMi/pkg/util"
+)
+
+// buildContainerAllocateResponse builds the allocate response for a single container.
+func (ps *PluginServer) buildContainerAllocateResponse(pod *v1.Pod, containerDevs device.ContainerDevices, rtInfoLookup map[string]RuntimeInfo) (*v1beta1.ContainerAllocateResponse, error) {
+	resp := &v1beta1.ContainerAllocateResponse{}
+
+	var (
+		IDs            []int32
+		memories       []*int64
+		cores          []*int32
+		ascendVNPUSpec string
+	)
+
+	for _, dev := range containerDevs {
+		d := ps.mgr.GetDeviceByUUID(dev.UUID)
+		if d == nil {
+			return nil, fmt.Errorf("unknown uuid: %s", dev.UUID)
+		}
+		IDs = append(IDs, d.PhyID)
+
+		if info, ok := rtInfoLookup[dev.UUID]; ok {
+			if ascendVNPUSpec == "" && info.Temp != "" {
+				ascendVNPUSpec = info.Temp
+			}
+			if info.Memory != nil {
+				memories = append(memories, info.Memory)
+			}
+			if info.Core != nil {
+				cores = append(cores, info.Core)
+			}
+		}
+	}
+
+	if len(IDs) == 0 {
+		return nil, fmt.Errorf("annotation %s value invalid", ps.allocAnno)
+	}
+	ascendVisibleDevices := fmt.Sprintf("%d", IDs[0])
+	for i := 1; i < len(IDs); i++ {
+		ascendVisibleDevices = fmt.Sprintf("%s,%d", ascendVisibleDevices, IDs[i])
+	}
+	resp.Envs = make(map[string]string)
+	resp.Envs["ASCEND_VISIBLE_DEVICES"] = ascendVisibleDevices
+
+	vnpuMode := pod.Annotations[VNPUModeAnnotation]
+	klog.V(4).Infof("Pod %s vnpu mode: %s", pod.Name, vnpuMode)
+	if vnpuMode == VNPUModeHamiCore {
+		// 1. Handle volume mount injection
+		var mounts []*v1beta1.Mount
+		// A.Huawei driver and SMI toolchain (Read-Only)
+		driverPaths := []string{
+			"/usr/local/bin/npu-smi",
+			"/etc/ascend_install.info",
+			"/usr/local/Ascend/driver/lib64/driver",
+			"/usr/local/Ascend/driver/version.info",
+		}
+		for _, p := range driverPaths {
+			mounts = append(mounts, &v1beta1.Mount{HostPath: p, ContainerPath: p, ReadOnly: true})
+		}
+
+		mounts = append(mounts, &v1beta1.Mount{
+			HostPath:      "/usr/local/hami-vnpu-core",
+			ContainerPath: "/hami-vnpu-core",
+			ReadOnly:      true,
+		})
+		// B. Inject HAMi library path by mounting /etc/ld.so.preload.
+		mounts = append(mounts, &v1beta1.Mount{
+			HostPath:      "/usr/local/hami-vnpu-core/ld.so.preload", // Template file on host
+			ContainerPath: "/etc/ld.so.preload",                      // Overwrites the target file in container
+			ReadOnly:      true,
+		})
+
+		// C. Shared directory for HAMi compute resource partitioning (Read/Write)
+		mounts = append(mounts, &v1beta1.Mount{
+			HostPath:      "/usr/local/hami-shared-region",
+			ContainerPath: "/hami-shared-region",
+			ReadOnly:      false,
+		})
+		resp.Mounts = mounts
+
+		// Set NPU_MEM_QUOTA
+		if len(memories) > 0 && memories[0] != nil {
+			resp.Envs["NPU_MEM_QUOTA"] = strconv.FormatInt(*memories[0], 10)
+			klog.V(4).InfoS("Memory quota set", "value", *memories[0])
+		}
+
+		// Set NPU_PRIORITY
+		if len(cores) > 0 && cores[0] != nil {
+			resp.Envs["NPU_PRIORITY"] = strconv.FormatInt(int64(*cores[0]), 10)
+			klog.V(4).InfoS("Core priority set", "value", *cores[0])
+		}
+
+		// Set GLOBAL_SHM_PATH based on the first device ID.
+		resp.Envs["NPU_GLOBAL_SHM_PATH"] = fmt.Sprintf("/hami-shared-region/%d_global_registry", IDs[0])
+		klog.V(5).Infof("Create %d_global_registry", IDs[0])
+	} else {
+		if ascendVNPUSpec != "" {
+			resp.Envs["ASCEND_VNPU_SPECS"] = ascendVNPUSpec
+		}
+	}
+	return resp, nil
+}
+
+// popNextContainerDevices finds and erases the first non-empty containerDevices
+// from podSingleDev. It mutates podSingleDev in place.
+func (ps *PluginServer) popNextContainerDevices(podSingleDev device.PodSingleDevice) (device.ContainerDevices, error) {
+	for i, ctrDevs := range podSingleDev {
+		if len(ctrDevs) > 0 {
+			podSingleDev[i] = device.ContainerDevices{}
+			return ctrDevs, nil
+		}
+	}
+	return nil, fmt.Errorf("no pending device allocation found")
+}
+
+// decodeDeviceAnnotations decodes the pod's device allocation annotation
+// (registered as hami.io/<commonword>-devices-to-allocate in InRequestDevices)
+// into a PodSingleDevice.
+func (ps *PluginServer) decodeDeviceAnnotations(pod *v1.Pod) (device.PodSingleDevice, error) {
+	pdevices, err := device.DecodePodDevices(device.InRequestDevices, pod.Annotations)
+	if err != nil {
+		return nil, err
+	}
+	pd, ok := pdevices[ps.commonWord]
+	if !ok {
+		return nil, fmt.Errorf("device %s not found in pod annotations", ps.commonWord)
+	}
+	return pd, nil
+}
+
+// buildRuntimeInfoLookup builds a UUID-to-RuntimeInfo lookup from the pod's allocAnno annotation.
+func (ps *PluginServer) buildRuntimeInfoLookup(pod *v1.Pod) (map[string]RuntimeInfo, error) {
+	anno, ok := pod.Annotations[ps.allocAnno]
+	if !ok {
+		return nil, fmt.Errorf("annotation %s not set", ps.allocAnno)
+	}
+	var rtInfo []RuntimeInfo
+	if err := json.Unmarshal([]byte(anno), &rtInfo); err != nil {
+		return nil, fmt.Errorf("annotation %s value %s invalid: %w", ps.allocAnno, anno, err)
+	}
+	lookup := make(map[string]RuntimeInfo, len(rtInfo))
+	for _, info := range rtInfo {
+		if info.UUID != "" {
+			lookup[info.UUID] = info
+		}
+	}
+	return lookup, nil
+}
+
+// patchErasedAnnotation patches the pod's device annotation with the given
+// podSingleDev. It also updates pod.Annotations in place.
+func (ps *PluginServer) patchErasedAnnotation(pod *v1.Pod, podSingleDev device.PodSingleDevice) error {
+	klog.V(5).Infof("After erase annotation, remaining devices: %v", podSingleDev)
+	newAnnoValue := device.EncodePodSingleDevice(podSingleDev)
+	newAnnos := map[string]string{
+		ps.toAllocDeviceAnno: newAnnoValue,
+	}
+	if err := util.PatchPodAnnotations(pod, newAnnos); err != nil {
+		return err
+	}
+	pod.Annotations[ps.toAllocDeviceAnno] = newAnnoValue
+	return nil
+}
+
+// podAllocationTrySuccess checks if all containers of this pod have been
+// allocated. If so, it sets bind-phase to "success" and releases the node
+// lock; otherwise it returns without setting bind-phase or releasing the lock,
+// waiting for the next Allocate call.
+func (ps *PluginServer) podAllocationTrySuccess(pod *v1.Pod) {
+	plugin.PodAllocationTrySuccess(ps.nodeName, ps.commonWord, NodeLockAscend, pod)
+}
+
+// podAllocationFailed sets bind-phase to "failed" and releases the node lock.
+func (ps *PluginServer) podAllocationFailed(pod *v1.Pod) {
+	plugin.PodAllocationFailed(ps.nodeName, pod, NodeLockAscend)
+}

internal/server/register.go

@@ -0,0 +1,152 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	"net"
+	"path"
+	"strings"
+	"time"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials/insecure"
+	"k8s.io/klog/v2"
+	"k8s.io/kubelet/pkg/apis/deviceplugin/v1beta1"
+
+	"github.com/Project-HAMi/HAMi/pkg/device"
+	"github.com/Project-HAMi/HAMi/pkg/util"
+)
+
+func (ps *PluginServer) watchAndRegister() {
+	timer := time.After(1 * time.Second)
+	for {
+		select {
+		case <-ps.stopCh:
+			klog.Infof("stop watch and register")
+			return
+		case <-timer:
+		}
+		unhealthy := ps.mgr.GetUnHealthIDs()
+		if len(unhealthy) > 0 {
+			if err := ps.mgr.UpdateDevice(); err != nil {
+				klog.Errorf("update device error: %v", err)
+				timer = time.After(5 * time.Second)
+				continue
+			}
+			ps.healthCh <- unhealthy[0]
+		}
+		err := ps.registerHAMi()
+		if err != nil {
+			klog.Errorf("register HAMi error: %v", err)
+			timer = time.After(5 * time.Second)
+		} else {
+			klog.V(3).Infof("register HAMi success")
+			timer = time.After(30 * time.Second)
+		}
+	}
+}
+
+func (ps *PluginServer) registerHAMi() error {
+	devs := ps.mgr.GetDevices()
+	apiDevices := make([]*device.DeviceInfo, 0, len(devs))
+	// hami currently believes that the index starts from 0 and is continuous.
+	for i, dev := range devs {
+		device := &device.DeviceInfo{
+			Index:   uint(i),
+			ID:      dev.UUID,
+			Count:   int32(ps.mgr.VDeviceCount()),
+			Devmem:  int32(dev.Memory),
+			Devcore: dev.AICore,
+			Type:    ps.mgr.CommonWord(),
+			Numa:    0,
+			Health:  dev.Health,
+		}
+		if strings.HasPrefix(device.Type, Ascend910Prefix) {
+			NetworkID, err := ps.getDeviceNetworkID(i, device.Type)
+			if err != nil {
+				return fmt.Errorf("get networkID error: %w", err)
+			}
+			device.CustomInfo = map[string]any{
+				"NetworkID": NetworkID,
+			}
+		}
+		apiDevices = append(apiDevices, device)
+	}
+	annos := make(map[string]string)
+	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")
+
+	if ps.mgr.IsHamiVnpuCore() {
+		annos[VNPUNodeSelectorAnnotation] = "true"
+		klog.V(4).Infof("Node %s has HamiVnpuCore enabled, patching annotation %s: true", ps.nodeName, VNPUNodeSelectorAnnotation)
+	} else {
+		annos[VNPUNodeSelectorAnnotation] = "false"
+	}
+
+	node, err := util.GetNode(ps.nodeName)
+	if err != nil {
+		return fmt.Errorf("get node %s error: %w", ps.nodeName, err)
+	}
+	err = util.PatchNodeAnnotations(node, annos)
+	if err != nil {
+		return fmt.Errorf("patch node %s annotations error: %w", ps.nodeName, err)
+	}
+	klog.V(5).Infof("patch node %s annotations: %v", ps.nodeName, annos)
+	return nil
+}
+
+func (ps *PluginServer) getDeviceNetworkID(idx int, deviceType string) (int, error) {
+	// For Ascend910C devices, all modules (dies) are interconnected via HCCS
+	if deviceType == Ascend910CType {
+		return 0, nil
+	}
+
+	if idx > 3 {
+		return 1, nil
+	}
+
+	return 0, nil
+}
+
+func (ps *PluginServer) registerKubelet() error {
+	conn, err := ps.dial(v1beta1.KubeletSocket, 5*time.Second)
+	if err != nil {
+		return err
+	}
+	defer func(conn *grpc.ClientConn) {
+		_ = conn.Close()
+	}(conn)
+	client := v1beta1.NewRegistrationClient(conn)
+	reqt := &v1beta1.RegisterRequest{
+		Version:      v1beta1.Version,
+		Endpoint:     path.Base(ps.socket),
+		ResourceName: ps.mgr.ResourceName(),
+		Options: &v1beta1.DevicePluginOptions{
+			GetPreferredAllocationAvailable: false,
+		},
+	}
+
+	_, err = client.Register(context.Background(), reqt)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (ps *PluginServer) dial(unixSocketPath string, timeout time.Duration) (*grpc.ClientConn, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), timeout)
+	defer cancel()
+	c, err := grpc.DialContext(ctx, unixSocketPath,
+		grpc.WithTransportCredentials(insecure.NewCredentials()),
+		grpc.WithBlock(),
+		grpc.WithContextDialer(func(ctx2 context.Context, addr string) (net.Conn, error) {
+			var d net.Dialer
+			return d.DialContext(ctx2, "unix", addr)
+		}),
+	)
+
+	if err != nil {
+		return nil, err
+	}
+	return c, nil
+}