package database import ( "github.com/zeromicro/go-zero/core/logx" clientCore "gitlink.org.cn/JointCloud/pcm-coordinator/api/client" "gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/scheduler/schedulers/option" "gitlink.org.cn/JointCloud/pcm-coordinator/api/internal/types" "gitlink.org.cn/JointCloud/pcm-coordinator/pkg/constants" "gitlink.org.cn/JointCloud/pcm-coordinator/pkg/models" "gitlink.org.cn/JointCloud/pcm-coordinator/pkg/tracker" "gorm.io/gorm" "strconv" "time" ) type AiStorage struct { DbEngin *gorm.DB } func (s *AiStorage) GetParticipants() (*types.ClusterListResp, error) { var resp types.ClusterListResp tx := s.DbEngin.Raw("select * from t_cluster where `deleted_at` IS NULL ORDER BY create_time Desc").Scan(&resp.List) if tx.Error != nil { logx.Errorf(tx.Error.Error()) return nil, tx.Error } return &resp, nil } func (s *AiStorage) GetClustersByAdapterId(id string) (*types.ClusterListResp, error) { var resp types.ClusterListResp tx := s.DbEngin.Raw("select * from t_cluster where `deleted_at` IS NULL and `adapter_id` = ? ORDER BY create_time Desc", id).Scan(&resp.List) if tx.Error != nil { logx.Errorf(tx.Error.Error()) return nil, tx.Error } return &resp, nil } func (s *AiStorage) GetClusterNameById(id string) (string, error) { var name string tx := s.DbEngin.Raw("select `description` from t_cluster where `id` = ?", id).Scan(&name) if tx.Error != nil { logx.Errorf(tx.Error.Error()) return "", tx.Error } return name, nil } func (s *AiStorage) GetAdapterNameById(id string) (string, error) { var name string tx := s.DbEngin.Raw("select `name` from t_adapter where `id` = ?", id).Scan(&name) if tx.Error != nil { logx.Errorf(tx.Error.Error()) return "", tx.Error } return name, nil } func (s *AiStorage) GetAdapterIdsByType(adapterType string) ([]string, error) { var list []types.AdapterInfo var ids []string db := s.DbEngin.Model(&types.AdapterInfo{}).Table("t_adapter") db = db.Where("type = ?", adapterType) err := db.Order("create_time desc").Find(&list).Error if err != nil { return nil, err } for _, info := range list { ids = append(ids, info.Id) } return ids, nil } func (s *AiStorage) GetAdaptersByType(adapterType string) ([]*types.AdapterInfo, error) { var list []*types.AdapterInfo db := s.DbEngin.Model(&types.AdapterInfo{}).Table("t_adapter") db = db.Where("type = ?", adapterType) err := db.Order("create_time desc").Find(&list).Error if err != nil { return nil, err } return list, nil } func (s *AiStorage) GetAiTasksByAdapterId(adapterId string) ([]*models.TaskAi, error) { var resp []*models.TaskAi db := s.DbEngin.Model(&models.TaskAi{}).Table("task_ai") db = db.Where("adapter_id = ?", adapterId) err := db.Order("commit_time desc").Find(&resp).Error if err != nil { return nil, err } return resp, nil } func (s *AiStorage) SaveTask(name string, strategyCode int64, synergyStatus int64) (int64, error) { // 构建主任务结构体 taskModel := models.Task{ Status: constants.Saved, Description: "ai task", Name: name, SynergyStatus: synergyStatus, Strategy: strategyCode, AdapterTypeDict: 1, CommitTime: time.Now(), } // 保存任务数据到数据库 tx := s.DbEngin.Create(&taskModel) if tx.Error != nil { return 0, tx.Error } return taskModel.Id, nil } func (s *AiStorage) SaveAiTask(taskId int64, option *option.AiOption, adapterName string, clusterId string, clusterName string, jobId string, status string, msg string) error { // 构建主任务结构体 aId, err := strconv.ParseInt(option.AdapterId, 10, 64) if err != nil { return err } cId, err := strconv.ParseInt(clusterId, 10, 64) if err != nil { return err } aiTaskModel := models.TaskAi{ TaskId: taskId, AdapterId: aId, AdapterName: adapterName, ClusterId: cId, ClusterName: clusterName, Name: option.TaskName, Replica: int64(option.Replica), JobId: jobId, TaskType: option.TaskType, Strategy: option.StrategyName, Status: status, Msg: msg, Card: option.ComputeCard, CommitTime: time.Now(), } // 保存任务数据到数据库 tx := s.DbEngin.Create(&aiTaskModel) if tx.Error != nil { return tx.Error } return nil } func (s *AiStorage) SaveClusterTaskQueue(adapterId string, clusterId string, queueNum int64) error { aId, err := strconv.ParseInt(adapterId, 10, 64) if err != nil { return err } cId, err := strconv.ParseInt(clusterId, 10, 64) if err != nil { return err } taskQueue := models.TClusterTaskQueue{ AdapterId: aId, ClusterId: cId, QueueNum: queueNum, } tx := s.DbEngin.Create(&taskQueue) if tx.Error != nil { return tx.Error } return nil } func (s *AiStorage) GetClusterTaskQueues(adapterId string, clusterId string) ([]*models.TClusterTaskQueue, error) { var taskQueues []*models.TClusterTaskQueue tx := s.DbEngin.Raw("select * from t_cluster_task_queue where `adapter_id` = ? and `cluster_id` = ?", adapterId, clusterId).Scan(&taskQueues) if tx.Error != nil { logx.Errorf(tx.Error.Error()) return nil, tx.Error } return taskQueues, nil } func (s *AiStorage) GetAiTaskIdByClusterIdAndTaskId(clusterId string, taskId string) (string, error) { var aiTask models.TaskAi tx := s.DbEngin.Raw("select * from task_ai where `cluster_id` = ? and `task_id` = ?", clusterId, taskId).Scan(&aiTask) if tx.Error != nil { logx.Errorf(tx.Error.Error()) return "", tx.Error } return aiTask.JobId, nil } func (s *AiStorage) GetClusterResourcesById(clusterId string) (*models.TClusterResource, error) { var clusterResource models.TClusterResource tx := s.DbEngin.Raw("select * from t_cluster_resource where `cluster_id` = ?", clusterId).Scan(&clusterResource) if tx.Error != nil { logx.Errorf(tx.Error.Error()) return nil, tx.Error } return &clusterResource, nil } func (s *AiStorage) SaveClusterResources(adapterId string, clusterId string, clusterName string, clusterType int64, cpuAvail float64, cpuTotal float64, memAvail float64, memTotal float64, diskAvail float64, diskTotal float64, gpuAvail float64, gpuTotal float64, cardTotal int64, topsTotal float64) error { cId, err := strconv.ParseInt(clusterId, 10, 64) if err != nil { return err } aId, err := strconv.ParseInt(adapterId, 10, 64) if err != nil { return err } clusterResource := models.TClusterResource{ AdapterId: aId, ClusterId: cId, ClusterName: clusterName, ClusterType: clusterType, CpuAvail: cpuAvail, CpuTotal: cpuTotal, MemAvail: memAvail, MemTotal: memTotal, DiskAvail: diskAvail, DiskTotal: diskTotal, GpuAvail: gpuAvail, GpuTotal: gpuTotal, CardTotal: cardTotal, CardTopsTotal: topsTotal, } tx := s.DbEngin.Create(&clusterResource) if tx.Error != nil { return tx.Error } // prometheus param := tracker.ClusterLoadRecord{ AdapterId: aId, ClusterName: clusterName, CpuAvail: cpuAvail, CpuTotal: cpuTotal, CpuUtilisation: clusterResource.CpuAvail / clusterResource.CpuTotal, MemoryAvail: memAvail, MemoryTotal: memTotal, MemoryUtilisation: clusterResource.MemAvail / clusterResource.MemTotal, DiskAvail: diskAvail, DiskTotal: diskTotal, DiskUtilisation: clusterResource.DiskAvail / clusterResource.DiskTotal, } tracker.SyncClusterLoad(param) return nil } func (s *AiStorage) UpdateClusterResources(clusterResource *models.TClusterResource) error { tx := s.DbEngin.Where("cluster_id = ?", clusterResource.ClusterId).Updates(clusterResource) if tx.Error != nil { return tx.Error } // prometheus param := tracker.ClusterLoadRecord{ AdapterId: clusterResource.AdapterId, ClusterName: clusterResource.ClusterName, CpuAvail: clusterResource.CpuAvail, CpuTotal: clusterResource.CpuTotal, CpuUtilisation: clusterResource.CpuAvail / clusterResource.CpuTotal, MemoryAvail: clusterResource.MemAvail, MemoryTotal: clusterResource.MemTotal, MemoryUtilisation: clusterResource.MemAvail / clusterResource.MemTotal, DiskAvail: clusterResource.DiskAvail, DiskTotal: clusterResource.DiskTotal, DiskUtilisation: clusterResource.DiskAvail / clusterResource.DiskTotal, } tracker.SyncClusterLoad(param) return nil } func (s *AiStorage) UpdateAiTask(task *models.TaskAi) error { tx := s.DbEngin.Updates(task) if tx.Error != nil { return tx.Error } return nil } func (s *AiStorage) GetStrategyCode(name string) (int64, error) { var strategy int64 sqlStr := `select t_dict_item.item_value from t_dict left join t_dict_item on t_dict.id = t_dict_item.dict_id where item_text = ? and t_dict.dict_code = 'schedule_Strategy'` //查询调度策略 err := s.DbEngin.Raw(sqlStr, name).Scan(&strategy).Error if err != nil { return strategy, nil } return strategy, nil } func (s *AiStorage) AddNoticeInfo(adapterId string, adapterName string, clusterId string, clusterName string, taskName string, noticeType string, incident string) { aId, err := strconv.ParseInt(adapterId, 10, 64) if err != nil { return } cId, err := strconv.ParseInt(clusterId, 10, 64) if err != nil { return } noticeInfo := clientCore.NoticeInfo{ AdapterId: aId, AdapterName: adapterName, ClusterId: cId, ClusterName: clusterName, NoticeType: noticeType, TaskName: taskName, Incident: incident, CreatedTime: time.Now(), } result := s.DbEngin.Table("t_notice").Create(¬iceInfo) if result.Error != nil { logx.Errorf("Task creation failure, err: %v", result.Error) } }