aiforge/modules/urfs_client/urchin/schedule.go

package urchin

import (
	"code.gitea.io/gitea/modules/labelmsg"
	"code.gitea.io/gitea/modules/setting"
	"encoding/json"
	"fmt"
	"strings"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/log"
)

type DecompressReq struct {
	SourceFile string `json:"source_file"`
	DestPath   string `json:"dest_path"`
}

const (
	modelSuffix = "models.zip"
)

var urfsClient Urchinfs

func getUrfsClient() {
	if urfsClient != nil {
		return
	}

	urfsClient = New()
}

func GetBackNpuModel(cloudbrainID int64, endpoint, bucket, objectKey, destPeerHost string) error {
	getUrfsClient()
	res, err := urfsClient.ScheduleDataToPeerByKey(endpoint, bucket, objectKey, destPeerHost)
	if err != nil {
		log.Error("ScheduleDataToPeerByKey failed:%v", err)
		return err
	}

	_, err = models.InsertScheduleRecord(&models.ScheduleRecord{
		CloudbrainID: cloudbrainID,
		EndPoint:     res.DataEndpoint,
		Bucket:       res.DataRoot,
		ObjectKey:    res.DataPath,
		ProxyServer:  destPeerHost,
		Status:       res.StatusCode,
	})
	if err != nil {
		log.Error("InsertScheduleRecord failed:%v", err)
		return err
	}

	switch res.StatusCode {
	case models.StorageScheduleSucceed:
		log.Info("ScheduleDataToPeerByKey succeed")
		decompress(res.DataRoot+"/"+res.DataPath, setting.Bucket+"/"+strings.TrimSuffix(res.DataPath, modelSuffix))
	case models.StorageScheduleProcessing:
		log.Info("ScheduleDataToPeerByKey processing")
	case models.StorageScheduleFailed:
		log.Error("ScheduleDataToPeerByKey failed:%s", res.StatusMsg)
		return fmt.Errorf("GetBackNpuModel failed:%s", res.StatusMsg)
	default:
		log.Info("ScheduleDataToPeerByKey failed, unknown StatusCode:%d", res.StatusCode)
		return fmt.Errorf("GetBackNpuModel failed, unknow StatusCode:%d", res.StatusCode)
	}

	return nil
}

func HandleScheduleRecords() error {
	getUrfsClient()
	records, err := models.GetSchedulingRecord()
	if err != nil {
		log.Error("GetSchedulingRecord failed:%v", err)
		return err
	}

	for _, record := range records {
		res, err := urfsClient.CheckScheduleTaskStatusByKey(record.EndPoint, record.Bucket, record.ObjectKey, record.ProxyServer)
		if err != nil {
			log.Error("CheckScheduleTaskStatusByKey(%d) failed:%v", record.ID, err)
			continue
		}

		record.Status = res.StatusCode
		models.UpdateScheduleCols(record, "status")

		switch res.StatusCode {
		case models.StorageScheduleSucceed:
			log.Info("ScheduleDataToPeerByKey(%s) succeed", record.ObjectKey)
			decompress(record.Bucket+"/"+record.ObjectKey, setting.Bucket+"/"+strings.TrimSuffix(record.ObjectKey, modelSuffix))
		case models.StorageScheduleProcessing:
			log.Info("ScheduleDataToPeerByKey(%s) processing", record.ObjectKey)
		case models.StorageScheduleFailed:
			log.Error("ScheduleDataToPeerByKey(%s) failed:%s", record.ObjectKey, res.StatusMsg)

		default:
			log.Info("ScheduleDataToPeerByKey(%s) failed, unknown StatusCode:%d", record.ObjectKey, res.StatusCode)
		}

	}

	return nil
}

func decompress(sourceFile, destPath string) {
	req, _ := json.Marshal(DecompressReq{
		SourceFile: sourceFile,
		DestPath:   destPath,
	})
	err := labelmsg.SendDecompressAttachToLabelOBS(string(req))
	if err != nil {
		log.Error("SendDecompressTask to labelsystem (%s) failed:%s", sourceFile, err.Error())
	}
}