msg_resp.go

package nm

import (
	"example.com/m/conf"
	"example.com/m/largeModel"
	"example.com/m/log"
	"example.com/m/models"
	"example.com/m/utils"
	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/crypto"
	nodemanagerV2 "github.com/odysseus/odysseus-protocol/gen/proto/go/nodemanager/v2"
	"math/big"
	"strconv"
	"time"
)

type WorkerMsgHandler func(params ...interface{}) *nodemanagerV2.WorkerMessage

type RespMsgHandler struct {
	nodeManager  *models.NodeManagerClient
	workerClient nodemanagerV2.NodeManagerService_RegisterWorkerClient
	handler      WorkerMsgHandler
	params       []interface{}
}

type RespMsgWorker struct {
	MsgPool chan *RespMsgHandler
}

func NewMsgRespWorker() *RespMsgWorker {
	return &RespMsgWorker{
		MsgPool: make(chan *RespMsgHandler, 0),
	}
}

func (o *RespMsgWorker) RegisterMsgResp(nodeManager *models.NodeManagerClient, workerClient nodemanagerV2.NodeManagerService_RegisterWorkerClient, handler WorkerMsgHandler, params []interface{}) {
	o.MsgPool <- &RespMsgHandler{
		nodeManager:  nodeManager,
		workerClient: workerClient,
		handler:      handler,
		params:       params,
	}
	log.Info("----------------add msg response-------------")
}

func (o *RespMsgWorker) SendMsgWorker() {
	for {
		select {
		case pool := <-o.MsgPool:
			{
				workerMsg := pool.handler(pool.params...)
				if workerMsg == nil {
					log.Warn("Send to node manager workerMsg is nil")
					continue
				}
				err := pool.workerClient.SendMsg(workerMsg)
				if err != nil {
					log.Error("Send msg to nm client failed:", err)
					return
				}
				log.Info("Worker client send message successfully")
			}
		}
	}
}

func HeartbeatResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Heartbeat response received params: ", params)
	serverTimestamp := params[0].(uint64)
	heartRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_HeartbeatResponse{
			HeartbeatResponse: &nodemanagerV2.HeartbeatResponse{
				Timestamp: serverTimestamp,
			},
		},
	}
	log.Info("---------------------------------------Send heart beat msg ------------------------------------")
	return heartRes
}

func RegisterInfoResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Register info response received params:", params)
	info := &nodemanagerV2.NodeInfo{
		MinerPubkey:    conf.GetConfig().SignPub,
		BenefitAddress: conf.GetConfig().BenefitAddress,
	}
	hardwareInfo := GetHardwareInfo()
	modelsInfo := params[0].(*largeModel.ModelHandler)
	readModels, err := modelsInfo.ScanModelsResp()
	if err != nil {
		log.WithError(err).Error("Error scanning models response failed")
		return nil
	}
	nowTimeStamp := time.Now().Unix()
	nowTimeBytes := big.NewInt(nowTimeStamp).Bytes()
	signHash := crypto.Keccak256Hash([]byte(info.String()),
		[]byte(hardwareInfo.String()),
		[]byte(readModels.String()),
		nowTimeBytes)
	log.WithField("hash", signHash.String()).Info("register message sign result")
	sign, _ := crypto.Sign(signHash.Bytes(), conf.GetConfig().SignPrivateKey)
	log.Info("register message sign:", common.Bytes2Hex(sign))
	nodeInfoRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_RegisteMessage{
			RegisteMessage: &nodemanagerV2.RegisteMessage{
				Info:            info,
				Hardware:        hardwareInfo,
				Models:          readModels,
				Timestamp:       nowTimeStamp,
				DeviceSignature: sign,
			},
		},
	}
	log.Info("---------------------------------------Send register info msg ------------------------------------")
	return nodeInfoRes
}

func NodeInfoResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Node info response received params:", params)
	hardwareInfo := GetHardwareInfo()
	modelsInfo := params[0].(*largeModel.ModelHandler)
	readModels, err := modelsInfo.GetRpcModelsResp()
	if err != nil {
		log.WithError(err).Error("Error getting rpc models response")
		return nil
	}
	nodeInfoRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_NodeInfo{
			NodeInfo: &nodemanagerV2.NodeInfoResponse{
				Info: &nodemanagerV2.NodeInfo{
					MinerPubkey:    conf.GetConfig().SignPub,
					BenefitAddress: conf.GetConfig().BenefitAddress,
				},
				Hardware: hardwareInfo,
				Models:   readModels,
			},
		},
	}
	log.Info("---------------------------------------Send node info msg ------------------------------------")
	return nodeInfoRes
}

func DeviceInfoResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Device info response received params:", params)
	hardwareInfo := GetHardwareInfo()
	deviceInfoRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_DeviceInfo{
			DeviceInfo: &nodemanagerV2.DeviceInfoMessage{
				Hardware:        hardwareInfo,
				DeviceSignature: []byte(""),
			},
		},
	}
	log.Info("---------------------------------------Send device info msg ------------------------------------")
	return deviceInfoRes
}

func DeviceUsageResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("DeviceUsageResp params :", params)
	hardwareInfo := GetHardwareInfo()
	ramUsage := int32((1 - float64(hardwareInfo.RAM.Free)/float64(hardwareInfo.RAM.Total)) * 100)
	diskUsage := int32((1 - float64(hardwareInfo.DISK.Free)/float64(hardwareInfo.DISK.Total)) * 100)
	deviceInfoRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_DeviceUsage{
			DeviceUsage: &nodemanagerV2.DeviceUsageResponse{
				Usage: &nodemanagerV2.HardwareUsage{
					CpuUsage:     hardwareInfo.CPU.Usage,
					RamUsage:     ramUsage,
					DiskUsage:    diskUsage,
					NetBandwidth: hardwareInfo.NET.Bandwidth,
				},
			},
		},
	}
	log.Info("---------------------------------------Send device usage msg ------------------------------------")
	return deviceInfoRes
}

func GpuUsageResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("DeviceUsageResp params :", params)
	hardwareInfo := GetHardwareInfo()
	if hardwareInfo == nil {
		return nil
	}
	gpusUsage := make([]*nodemanagerV2.GPUUsage, 0)
	for _, gpuInfo := range hardwareInfo.GPU {
		usage := &nodemanagerV2.GPUUsage{
			Seq:     gpuInfo.Seq,
			MemFree: gpuInfo.MemFree,
			Usage:   gpuInfo.Usage,
			Temp:    gpuInfo.Temp,
			PowerRt: gpuInfo.PowerRt,
		}
		gpusUsage = append(gpusUsage, usage)
	}
	deviceInfoRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_GpuUsage{
			GpuUsage: &nodemanagerV2.GPUUsageResponse{
				Usages: gpusUsage,
			},
		},
	}
	log.Info("---------------------------------------Send gpu usage msg ------------------------------------")
	return deviceInfoRes
}

func ModelListResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("DeviceUsageResp params :", params)
	modelsInfo := params[0].(*largeModel.ModelHandler)
	readModels, err := modelsInfo.GetRpcModelsResp()
	if err != nil {
		return nil
	}
	modelListInfoRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_ModelsInfo{
			ModelsInfo: readModels,
		},
	}
	log.Info("---------------------------------------Send model list msg ------------------------------------")
	return modelListInfoRes
}

func SubmitResultResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	//log.Info("Handler task submit result resp received params:", params)
	taskId := params[0].(string)
	containerSign := params[1].([]byte)
	minerSign := params[2].([]byte)
	taskExecResult := params[3].(*models.TaskResult)
	isSuccess := params[4].(bool)
	n := &nodemanagerV2.SubmitTaskResult{
		TaskId:              taskId,
		ContainerSignature:  containerSign,
		MinerSignature:      minerSign,
		TaskExecuteCode:     taskExecResult.TaskHttpStatusCode,
		TaskExecuteError:    taskExecResult.TaskExecError,
		TaskResultHeader:    taskExecResult.TaskHttpHeaders,
		TaskExecuteDuration: uint64(taskExecResult.TaskExecTime),
		IsSuccessed:         isSuccess,
		TaskResultBody:      taskExecResult.TaskRespBody,
	}
	submitResultMsgRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_SubmitTaskResult{
			SubmitTaskResult: n,
		},
	}
	log.Info("---------------------------------------Send task result msg ------------------------------------")
	return submitResultMsgRes
}

func FetchStandardTaskResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	//log.Info("Handler task submit result resp received params:", params)
	fetchStandardTaskMsgRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_FetchStandardTask{
			FetchStandardTask: &nodemanagerV2.FetchStandardTask{
				TaskType: 998,
			},
		},
	}
	log.Info("---------------------------------------Send fetch standard task msg ------------------------------------")
	return fetchStandardTaskMsgRes
}

func RespTaskAck(params ...interface{}) *nodemanagerV2.WorkerMessage {
	taskId := params[0].(string)
	canExecute := params[1].(bool)
	bootUpTime := params[2].(int64)
	queueWaitTime := params[3].(int64)
	executeTime := params[4].(int64)
	taskAckMsgRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_SubmitTaskAck{
			SubmitTaskAck: &nodemanagerV2.SubmitTaskAck{
				TaskId:        taskId,
				CanExecute:    canExecute,
				BootUpTime:    bootUpTime,
				QueueWaitTime: queueWaitTime,
				ExecuteTime:   executeTime,
			},
		},
	}
	log.WithField("taskId", taskId).Info("---------------------------------------Send task ack msg ------------------------------------")
	return taskAckMsgRes
}

func GoodbyeResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Goodbye resp received params:", params)
	reason := ""
	if len(params) > 0 {
		reason = params[0].(string)
	}
	goodbyeMsgRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_GoodbyeMessage{
			GoodbyeMessage: &nodemanagerV2.GoodbyeMessage{
				Reason: reason,
			},
		},
	}
	log.Info("---------------------------------------Send good bye msg ------------------------------------")
	return goodbyeMsgRes
}

func AddModelInstalledResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Add model installed info response received params:", params)
	installedModels := make([]*nodemanagerV2.InstalledModel, 0)
	for _, param := range params {
		model := param.(*nodemanagerV2.InstalledModel)
		installedModels = append(installedModels, model)
	}
	deviceInfoRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_AddModelInstalled{
			AddModelInstalled: &nodemanagerV2.AddModelInstalled{
				Models: installedModels,
			},
		},
	}
	log.Info("---------------------------------------Send add model installed info msg ------------------------------------")
	return deviceInfoRes
}

func DelModelInstalledResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Del model installed info response received params:", params)
	deviceInfoRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_DelModelInstalled{
			DelModelInstalled: &nodemanagerV2.DelModelInstalled{
				ModelIds: []string{params[0].(string)},
			},
		},
	}
	log.Info("---------------------------------------Del model installed info msg ------------------------------------")
	return deviceInfoRes
}

func AddModelRunningResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Add model running response received params:", params)
	runningModels := make([]*nodemanagerV2.RunningModel, 0)
	for _, param := range params {
		runningModel := param.(*nodemanagerV2.RunningModel)
		runningModels = append(runningModels, runningModel)
	}
	addModelRunningRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_AddModelRunning{
			AddModelRunning: &nodemanagerV2.AddModelRunning{
				Models: runningModels,
			},
		},
	}
	log.Info("---------------------------------------Send Add model running response msg ------------------------------------")
	return addModelRunningRes
}

func DelModelRunningResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Del model running response received params:", params)
	delModelRunningRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_DelModeRunning{
			DelModeRunning: &nodemanagerV2.DelModelRunning{
				ModelIds: []string{params[0].(string)},
			},
		},
	}
	log.Info("---------------------------------------Send del model running response msg ------------------------------------")
	return delModelRunningRes
}

func BenefitAddrUpdateResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Benefit addr update response received params:", params)
	addr := params[0].(string)
	benefitAddrUpdateResp := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_BenefitAddrUpdate{
			BenefitAddrUpdate: &nodemanagerV2.BenefitAddrUpdate{
				BenefitAddress: addr,
			},
		},
	}
	log.Info("---------------------------------------Send Benefit addr update response msg ------------------------------------")
	return benefitAddrUpdateResp
}

func RunningModelStatusResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Running Model Status response received params:", params)
	info := params[0].(*models.ModelInfo)
	runningModelStatusResp := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_RunningModelStatus{
			RunningModelStatus: &nodemanagerV2.RunningModelStatus{
				ModelId:       strconv.FormatUint(info.TaskId, 10),
				LastWorkTime:  info.LastWorkTime,
				TotalRunCount: info.TotalRunCount,
				ExecTime:      info.EstimatExeTime,
			},
		},
	}
	log.Info("---------------------------------------Send running model status response msg ------------------------------------")
	return runningModelStatusResp
}

func InstallModelStatusResp(params ...interface{}) *nodemanagerV2.WorkerMessage {
	log.Info("Install Model Status response received params:", params)
	modelId := params[0].(string)
	lastRunTime := params[1].(int64)
	installModelStatusRes := &nodemanagerV2.WorkerMessage{
		Message: &nodemanagerV2.WorkerMessage_InstalledModelStatus{
			InstalledModelStatus: &nodemanagerV2.InstalledModelStatus{
				ModelId:     modelId,
				LastRunTime: lastRunTime,
			},
		},
	}
	log.Info("---------------------------------------Send install model status response msg ------------------------------------")
	return installModelStatusRes
}

func GetHardwareInfo() *nodemanagerV2.HardwareInfo {
	hardwareInfo := utils.GetHardwareInfo(conf.GetConfig().HardwareUrl)
	if hardwareInfo == nil {
		return nil
	}
	gpusInfo := make([]*nodemanagerV2.GPUInfo, 0)
	var diskTotal, diskFree int64
	for _, disk := range hardwareInfo.Data.Disk {
		for _, point := range disk.MountPoints {
			if point == "/" {
				diskTotal += disk.SizeBytes
				diskFree += disk.FreeBytes
			}
		}
	}
	diskTotal = diskTotal * conf.GetConfig().DiskUsage
	diskFree = diskFree * conf.GetConfig().DiskUsage
	var macAddr string
	var bandWidth int32
	for _, net := range hardwareInfo.Data.Networks {
		if net.Device == "docker0" {
			macAddr = net.Mac
			bandWidth = net.Speed
		}
	}
	for _, gpu := range hardwareInfo.Data.Gpus {
		gpuInfo := &nodemanagerV2.GPUInfo{
			Seq:         gpu.Seq,
			Uuid:        gpu.Uuid,
			Model:       gpu.Model,
			Performance: gpu.Performance,
			PowerRating: gpu.PowerRating,
			MemTotal:    gpu.MemTotal,
			MemFree:     gpu.MemFree,
			Usage:       gpu.Usage,
			Temp:        gpu.Temp,
			PowerRt:     gpu.PowerRt,
		}
		gpusInfo = append(gpusInfo, gpuInfo)
	}
	res := &nodemanagerV2.HardwareInfo{
		CPU: &nodemanagerV2.CPUInfo{
			Model:   hardwareInfo.Data.Cpus.Model,
			Number:  hardwareInfo.Data.Cpus.Number,
			Cores:   hardwareInfo.Data.Cpus.Cores,
			Threads: hardwareInfo.Data.Cpus.Threads,
			Usage:   hardwareInfo.Data.Cpus.Usage,
		},
		GPU: gpusInfo,
		RAM: &nodemanagerV2.MemoryInfo{
			Total: hardwareInfo.Data.Mem.Total,
			Free:  hardwareInfo.Data.Mem.Free,
		},
		DISK: &nodemanagerV2.DiskInfo{
			Total: diskTotal,
			Free:  diskFree,
		},
		NET: &nodemanagerV2.NetInfo{
			Ip:        conf.GetConfig().ExternalIp,
			Mac:       macAddr,
			Bandwidth: bandWidth,
		},
	}
	return res
}