402 lines
12 KiB
Go
402 lines
12 KiB
Go
|
package mq
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"bufio"
|
|||
|
|
"context"
|
|||
|
|
"encoding/base64"
|
|||
|
|
"encoding/json"
|
|||
|
|
"fmt"
|
|||
|
|
"git.hpds.cc/pavement/hpds_node"
|
|||
|
|
"github.com/emirpasic/gods/lists/arraylist"
|
|||
|
|
"github.com/fsnotify/fsnotify"
|
|||
|
|
"github.com/gammazero/workerpool"
|
|||
|
|
"github.com/minio/minio-go/v7"
|
|||
|
|
"github.com/minio/minio-go/v7/pkg/credentials"
|
|||
|
|
"github.com/shirou/gopsutil/v3/host"
|
|||
|
|
"go.uber.org/zap"
|
|||
|
|
"io"
|
|||
|
|
"net/http"
|
|||
|
|
"os"
|
|||
|
|
"os/exec"
|
|||
|
|
"path"
|
|||
|
|
"strconv"
|
|||
|
|
"strings"
|
|||
|
|
"sync"
|
|||
|
|
"taskExecute/config"
|
|||
|
|
"taskExecute/pkg/compress"
|
|||
|
|
"taskExecute/pkg/docker"
|
|||
|
|
"taskExecute/pkg/download"
|
|||
|
|
"taskExecute/pkg/utils"
|
|||
|
|
"taskExecute/proto"
|
|||
|
|
"time"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
var (
|
|||
|
|
wg sync.WaitGroup
|
|||
|
|
TaskList map[string]docker.ContainerStore
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
func TaskExecuteHandler(data []byte) (byte, []byte) {
|
|||
|
|
fmt.Println("接收数据", string(data))
|
|||
|
|
cmd := new(InstructionReq)
|
|||
|
|
err := json.Unmarshal(data, cmd)
|
|||
|
|
if err != nil {
|
|||
|
|
return 0x0B, []byte(err.Error())
|
|||
|
|
}
|
|||
|
|
switch cmd.Command {
|
|||
|
|
case TaskExecute:
|
|||
|
|
//任务执行
|
|||
|
|
waitWorkerStartFinish(config.WorkPool, cmd.Payload.(map[string]interface{}), ModelTaskExecuteHandler)
|
|||
|
|
case ModelIssueRepeater:
|
|||
|
|
//模型下发
|
|||
|
|
waitWorkerStartFinish(config.WorkPool, cmd.Payload.(map[string]interface{}), ModelIssueRepeaterHandler)
|
|||
|
|
}
|
|||
|
|
return byte(cmd.Command), nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func waitWorkerStartFinish(wp *workerpool.WorkerPool, payload map[string]interface{}, f func(payload map[string]interface{})) {
|
|||
|
|
startStop := make(chan time.Time, 2)
|
|||
|
|
wp.Submit(func() {
|
|||
|
|
startStop <- time.Now()
|
|||
|
|
f(payload)
|
|||
|
|
startStop <- time.Now()
|
|||
|
|
})
|
|||
|
|
fmt.Println("Task started at:", <-startStop)
|
|||
|
|
fmt.Println("Task finished at:", <-startStop)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// execCommand 执行命令
|
|||
|
|
func execCommandWait(commandName string, params []string) bool {
|
|||
|
|
cmd := exec.Command(commandName, params...)
|
|||
|
|
|
|||
|
|
//显示运行的命令
|
|||
|
|
fmt.Println(cmd.Args)
|
|||
|
|
|
|||
|
|
stdout, err := cmd.StdoutPipe()
|
|||
|
|
|
|||
|
|
if err != nil {
|
|||
|
|
fmt.Println(err)
|
|||
|
|
return false
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
_ = cmd.Start()
|
|||
|
|
|
|||
|
|
reader := bufio.NewReader(stdout)
|
|||
|
|
|
|||
|
|
//实时循环读取输出流中的一行内容
|
|||
|
|
for {
|
|||
|
|
wg.Add(1)
|
|||
|
|
line, err2 := reader.ReadString('\n')
|
|||
|
|
if err2 != nil || io.EOF == err2 {
|
|||
|
|
break
|
|||
|
|
}
|
|||
|
|
config.Logger.Info("执行命令",
|
|||
|
|
zap.String("execCommandWait", line))
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
_ = cmd.Wait()
|
|||
|
|
wg.Done()
|
|||
|
|
return true
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func ModelIssueRepeaterHandler(payload map[string]interface{}) {
|
|||
|
|
hi, _ := host.Info()
|
|||
|
|
if payload["nodeGuid"].(string) == hi.HostID {
|
|||
|
|
fileUrl := payload["dockerFile"].(string)
|
|||
|
|
modelVersion := payload["modelVersion"].(string)
|
|||
|
|
downFileName := path.Base(fileUrl)
|
|||
|
|
//判断文件后缀名
|
|||
|
|
fileType := path.Ext(downFileName)
|
|||
|
|
fileNameOnly := strings.TrimSuffix(downFileName, fileType)
|
|||
|
|
dFile := path.Join(config.Cfg.TmpPath, fileNameOnly, downFileName)
|
|||
|
|
//执行文件下载
|
|||
|
|
controller := download.ThreadController{
|
|||
|
|
ThreadCount: download.ThreadCount,
|
|||
|
|
FileUrl: fileUrl,
|
|||
|
|
DownloadFolder: dFile,
|
|||
|
|
DownloadFileName: downFileName,
|
|||
|
|
Logger: config.Logger,
|
|||
|
|
}
|
|||
|
|
controller.Download(download.OneThreadDownloadSize)
|
|||
|
|
if strings.ToLower(fileType) == ".zip" {
|
|||
|
|
err := compress.UnzipFromFile(path.Join(config.Cfg.TmpPath, fileNameOnly), dFile)
|
|||
|
|
if err != nil {
|
|||
|
|
controller.Logger.With(zap.String("文件解压缩", path.Join(config.Cfg.TmpPath, downFileName))).
|
|||
|
|
Error("发生错误", zap.Error(err))
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
dFile = path.Join(config.Cfg.TmpPath, fileNameOnly, fileNameOnly+".tar")
|
|||
|
|
}
|
|||
|
|
//docker 导入并运行
|
|||
|
|
imgName := fmt.Sprintf("%s:%s", fileNameOnly, modelVersion)
|
|||
|
|
if strings.ToLower(path.Ext(dFile)) == ".tar" {
|
|||
|
|
dCli := docker.NewDockerClient()
|
|||
|
|
err := dCli.ImportImage(imgName, "latest", dFile)
|
|||
|
|
//err = dCli.LoadImage(dFile)
|
|||
|
|
if err != nil {
|
|||
|
|
controller.Logger.With(zap.String("导入docker的文件", dFile)).
|
|||
|
|
Error("发生错误", zap.Error(err))
|
|||
|
|
}
|
|||
|
|
//设置data目录
|
|||
|
|
dataPath := path.Join(config.Cfg.TmpPath, fileNameOnly, "data")
|
|||
|
|
_ = os.MkdirAll(dataPath, os.ModePerm)
|
|||
|
|
vol := make(map[string]string)
|
|||
|
|
vol[path.Join(dataPath, payload["inPath"].(string))] = payload["inPath"].(string)
|
|||
|
|
vol[path.Join(dataPath, payload["outPath"].(string))] = payload["outPath"].(string)
|
|||
|
|
//docker运行
|
|||
|
|
modelCommand := strings.Split(payload["modelCommand"].(string), " ")
|
|||
|
|
|
|||
|
|
dstPort := dCli.CreateContainer(fileNameOnly, imgName, modelCommand, vol, strconv.Itoa(payload["mappedPort"].(int)))
|
|||
|
|
//保存到本地临时文件
|
|||
|
|
item := docker.ContainerStore{
|
|||
|
|
ModelId: payload["modelId"].(int64),
|
|||
|
|
NodeId: payload["nodeId"].(int64),
|
|||
|
|
Name: fileNameOnly,
|
|||
|
|
ImgName: imgName,
|
|||
|
|
Volumes: vol,
|
|||
|
|
SrcPort: strconv.Itoa(payload["mappedPort"].(int)),
|
|||
|
|
DstPort: dstPort,
|
|||
|
|
Command: modelCommand,
|
|||
|
|
HttpUrl: payload["httpUrl"].(string),
|
|||
|
|
}
|
|||
|
|
docker.ContainerList = append(docker.ContainerList, item)
|
|||
|
|
docker.SaveStore()
|
|||
|
|
cli := GetMqClient("task-response", 1)
|
|||
|
|
ap := cli.EndPoint.(hpds_node.AccessPoint)
|
|||
|
|
res := new(InstructionReq)
|
|||
|
|
res.Command = ModelIssueResponse
|
|||
|
|
res.Payload = item
|
|||
|
|
pData, _ := json.Marshal(res)
|
|||
|
|
_ = GenerateAndSendData(ap, pData)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func ModelTaskExecuteHandler(payload map[string]interface{}) {
|
|||
|
|
hi, _ := host.Info()
|
|||
|
|
if payload["nodeGuid"] == hi.HostID {
|
|||
|
|
if len(payload["subDataset"].(string)) > 0 {
|
|||
|
|
sf := hpds_node.NewStreamFunction(
|
|||
|
|
payload["subDataset"].(string),
|
|||
|
|
hpds_node.WithMqAddr(fmt.Sprintf("%s:%d", config.Cfg.Node.Host, config.Cfg.Node.Port)),
|
|||
|
|
hpds_node.WithObserveDataTags(payload["subDataTag"].(byte)),
|
|||
|
|
hpds_node.WithCredential(config.Cfg.Node.Token),
|
|||
|
|
)
|
|||
|
|
err := sf.Connect()
|
|||
|
|
must(config.Logger, err)
|
|||
|
|
nodeInfo := HpdsMqNode{
|
|||
|
|
MqType: 2,
|
|||
|
|
Topic: payload["subDataset"].(string),
|
|||
|
|
Node: config.Cfg.Node,
|
|||
|
|
EndPoint: sf,
|
|||
|
|
}
|
|||
|
|
_ = sf.SetHandler(func(data []byte) (byte, []byte) {
|
|||
|
|
|
|||
|
|
//查询docker是否已经开启
|
|||
|
|
issue := new(docker.ContainerStore)
|
|||
|
|
_ = json.Unmarshal([]byte(payload["issueResult"].(string)), issue)
|
|||
|
|
dCli := docker.NewDockerClient()
|
|||
|
|
cList, err := dCli.SearchImage(issue.Name)
|
|||
|
|
if err != nil {
|
|||
|
|
|
|||
|
|
}
|
|||
|
|
if len(cList) > 0 {
|
|||
|
|
if len(payload["workflow"].(string)) > 0 {
|
|||
|
|
//是否设置工作流程
|
|||
|
|
wf := new(Workflow)
|
|||
|
|
err = json.Unmarshal([]byte(payload["workflow"].(string)), wf)
|
|||
|
|
if err != nil {
|
|||
|
|
|
|||
|
|
}
|
|||
|
|
if len(payload["datasetPath"].(string)) > 0 {
|
|||
|
|
//数据集处理
|
|||
|
|
opt := &minio.Options{
|
|||
|
|
Creds: credentials.NewStaticV4(config.Cfg.Minio.AccessKeyId, config.Cfg.Minio.SecretAccessKey, ""),
|
|||
|
|
Secure: false,
|
|||
|
|
}
|
|||
|
|
cli, _ := minio.New(config.Cfg.Minio.Endpoint, opt)
|
|||
|
|
doneCh := make(chan struct{})
|
|||
|
|
defer close(doneCh)
|
|||
|
|
objectCh := cli.ListObjects(context.Background(), config.Cfg.Minio.Bucket, minio.ListObjectsOptions{
|
|||
|
|
Prefix: payload["datasetPath"].(string),
|
|||
|
|
Recursive: true,
|
|||
|
|
})
|
|||
|
|
for object := range objectCh {
|
|||
|
|
file, _ := cli.GetObject(context.Background(), config.Cfg.Minio.Bucket, object.Key, minio.GetObjectOptions{})
|
|||
|
|
imgByte, _ := io.ReadAll(file)
|
|||
|
|
|
|||
|
|
f := proto.FileCapture{
|
|||
|
|
FileName: object.Key,
|
|||
|
|
File: base64.StdEncoding.EncodeToString(imgByte),
|
|||
|
|
DatasetName: payload["datasetName"].(string),
|
|||
|
|
CaptureTime: object.LastModified.Unix(),
|
|||
|
|
}
|
|||
|
|
ProcessWorkflow(payload, f, wf)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
}
|
|||
|
|
} else {
|
|||
|
|
f := new(proto.FileCapture)
|
|||
|
|
err := json.Unmarshal(data, f)
|
|||
|
|
if err != nil {
|
|||
|
|
|
|||
|
|
}
|
|||
|
|
if len(f.File) > 0 {
|
|||
|
|
|
|||
|
|
i := strings.Index(f.File, ",")
|
|||
|
|
dec := base64.NewDecoder(base64.StdEncoding, strings.NewReader(f.File[i+1:]))
|
|||
|
|
if len(payload["httpUrl"].(string)) > 0 {
|
|||
|
|
_ = os.MkdirAll(path.Join(config.Cfg.TmpPath, payload["subDataset"].(string)), os.ModePerm)
|
|||
|
|
tmpFile, _ := os.Create(path.Join(config.Cfg.TmpPath, payload["subDataset"].(string), f.FileName))
|
|||
|
|
_, err = io.Copy(tmpFile, dec)
|
|||
|
|
|
|||
|
|
reqUrl := fmt.Sprintf("http://localhost:%s/%s", issue.DstPort, issue.HttpUrl)
|
|||
|
|
response, err := http.Post(reqUrl, "multipart/form-data", dec)
|
|||
|
|
if err != nil {
|
|||
|
|
config.Logger.With(zap.String("源文件名", f.FileName)).
|
|||
|
|
With(zap.String("临时文件名", path.Join(config.Cfg.TmpPath, payload["subDataset"].(string), f.FileName))).
|
|||
|
|
Error("文件提交", zap.Error(err))
|
|||
|
|
}
|
|||
|
|
defer func() {
|
|||
|
|
_ = response.Body.Close()
|
|||
|
|
config.Logger.With(zap.String("源文件名", f.FileName)).
|
|||
|
|
With(zap.String("临时文件名", path.Join(config.Cfg.TmpPath, payload["subDataset"].(string), f.FileName))).
|
|||
|
|
Info("模型识别")
|
|||
|
|
}()
|
|||
|
|
body, err := io.ReadAll(response.Body)
|
|||
|
|
if err != nil {
|
|||
|
|
config.Logger.With(zap.String("源文件名", f.FileName)).
|
|||
|
|
With(zap.String("临时文件名", path.Join(config.Cfg.TmpPath, payload["subDataset"].(string), f.FileName))).
|
|||
|
|
Error("模型识别", zap.Error(err))
|
|||
|
|
}
|
|||
|
|
cli := GetMqClient("task-response", 1)
|
|||
|
|
ap := cli.EndPoint.(hpds_node.AccessPoint)
|
|||
|
|
res := new(InstructionReq)
|
|||
|
|
res.Command = TaskResponse
|
|||
|
|
res.Payload = body
|
|||
|
|
pData, _ := json.Marshal(res)
|
|||
|
|
_ = GenerateAndSendData(ap, pData)
|
|||
|
|
}
|
|||
|
|
if len(payload["inPath"].(string)) > 0 {
|
|||
|
|
outPath := ""
|
|||
|
|
for k, v := range issue.Volumes {
|
|||
|
|
if v == payload["outPath"].(string) {
|
|||
|
|
outPath = k
|
|||
|
|
break
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
//创建一个监控对象
|
|||
|
|
watch, err := fsnotify.NewWatcher()
|
|||
|
|
if err != nil {
|
|||
|
|
config.Logger.Error("创建文件监控", zap.Error(err))
|
|||
|
|
}
|
|||
|
|
defer func(watch *fsnotify.Watcher) {
|
|||
|
|
_ = watch.Close()
|
|||
|
|
}(watch)
|
|||
|
|
|
|||
|
|
err = watch.Add(outPath)
|
|||
|
|
if err != nil {
|
|||
|
|
config.Logger.With(zap.String("监控目录", outPath)).
|
|||
|
|
Error("创建文件监控", zap.Error(err))
|
|||
|
|
}
|
|||
|
|
for k, v := range issue.Volumes {
|
|||
|
|
if v == payload["inPath"].(string) {
|
|||
|
|
_ = os.MkdirAll(k, os.ModePerm)
|
|||
|
|
tmpFile, _ := os.Create(path.Join(k, f.FileName))
|
|||
|
|
_, err = io.Copy(tmpFile, dec)
|
|||
|
|
break
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
list := arraylist.New() // empty
|
|||
|
|
t1 := time.NewTicker(1 * time.Second)
|
|||
|
|
go func() {
|
|||
|
|
for {
|
|||
|
|
select {
|
|||
|
|
case ev := <-watch.Events:
|
|||
|
|
{
|
|||
|
|
//判断事件发生的类型，如下5种
|
|||
|
|
// Create 创建
|
|||
|
|
// Write 写入
|
|||
|
|
// Remove 删除
|
|||
|
|
// Rename 重命名
|
|||
|
|
// Chmod 修改权限
|
|||
|
|
if ev.Op&fsnotify.Create == fsnotify.Create {
|
|||
|
|
config.Logger.Info("创建文件", zap.String("文件名", ev.Name))
|
|||
|
|
list.Add(ev.Name)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
case <-t1.C:
|
|||
|
|
{
|
|||
|
|
if list.Size() > 0 {
|
|||
|
|
returnFileHandleResult(list, payload, issue)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
case err = <-watch.Errors:
|
|||
|
|
{
|
|||
|
|
config.Logger.With(zap.String("监控目录", outPath)).
|
|||
|
|
Error("文件监控", zap.Error(err))
|
|||
|
|
return
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}()
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
return payload["subDataTag"].(byte), nil
|
|||
|
|
})
|
|||
|
|
MqList = append(MqList, nodeInfo)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
func returnFileHandleResult(list *arraylist.List, payload map[string]interface{}, issue *docker.ContainerStore) {
|
|||
|
|
var (
|
|||
|
|
mu sync.RWMutex
|
|||
|
|
wgp sync.WaitGroup
|
|||
|
|
resTime time.Duration
|
|||
|
|
)
|
|||
|
|
mu.Lock()
|
|||
|
|
defer mu.Unlock()
|
|||
|
|
startTime := time.Now()
|
|||
|
|
for i := 0; i < list.Size(); i++ {
|
|||
|
|
if fn, ok := list.Get(0); ok {
|
|||
|
|
if utils.PathExists(fn.(string)) {
|
|||
|
|
wgp.Add(1)
|
|||
|
|
go func() {
|
|||
|
|
mr := new(proto.ModelResult)
|
|||
|
|
src := utils.ReadFile(fn.(string))
|
|||
|
|
|
|||
|
|
if src != nil {
|
|||
|
|
mr.File = base64.StdEncoding.EncodeToString(src)
|
|||
|
|
mr.TaskCode = utils.GetFileName(fn.(string))
|
|||
|
|
mr.TaskId = int64(payload["taskId"].(float64))
|
|||
|
|
mr.FileName = utils.GetFileNameAndExt(fn.(string))
|
|||
|
|
mr.DatasetName = payload["datasetName"].(string)
|
|||
|
|
mr.SubDataset = payload["subDataset"].(string)
|
|||
|
|
mr.FileMd5 = utils.GetFileMd5(src)
|
|||
|
|
mr.ModelId = int64(payload["modelId"].(float64))
|
|||
|
|
mr.NodeId = int64(payload["nodeId"].(float64))
|
|||
|
|
mr.StartTime = startTime.Unix()
|
|||
|
|
mr.FinishTime = time.Now().Unix()
|
|||
|
|
cli := GetMqClient("task-response", 1)
|
|||
|
|
ap := cli.EndPoint.(hpds_node.AccessPoint)
|
|||
|
|
res := new(InstructionReq)
|
|||
|
|
res.Command = TaskResponse
|
|||
|
|
res.Payload = mr
|
|||
|
|
pData, _ := json.Marshal(res)
|
|||
|
|
_ = GenerateAndSendData(ap, pData)
|
|||
|
|
}
|
|||
|
|
wg.Done()
|
|||
|
|
}()
|
|||
|
|
wg.Wait()
|
|||
|
|
resTime = time.Since(startTime)
|
|||
|
|
config.Logger.Info("返回任务完成",
|
|||
|
|
zap.String("文件名", fn.(string)),
|
|||
|
|
zap.Duration("运行时间", resTime),
|
|||
|
|
)
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
}
|