eta_index_lib/logic/ai_predict_model/index.go

package ai_predict_model

import (
	"eta/eta_index_lib/models"
	"eta/eta_index_lib/models/ai_predict_model"
	"eta/eta_index_lib/utils"
	"fmt"
	"time"
)

// HandleTaskRecordFailByTaskRecord
// @Description: 任务标记失败
// @author: Roc
// @datetime 2025-05-09 16:24:48
// @param taskType string
// @param indexTaskRecordInfo *models.IndexTaskRecord
// @param indexConfigItem *ai_predict_model.AiPredictModelIndexConfig
// @param indexItem *ai_predict_model.AiPredictModelIndex
// @param errMsg string
func HandleTaskRecordFailByTaskRecord(taskType string, indexTaskRecordInfo *models.IndexTaskRecord, indexConfigItem *ai_predict_model.AiPredictModelIndexConfig, indexItem *ai_predict_model.AiPredictModelIndex, errMsg string) {
	var err error
	defer func() {
		if err != nil {
			utils.FileLog.Error(fmt.Sprintf(`HandleTaskRecordFailByTaskRecord err:%v`, err))
		}
	}()

	// 修改子任务状态
	indexTaskRecordInfo.Status = `处理失败`
	indexTaskRecordInfo.Remark = errMsg
	indexTaskRecordInfo.ModifyTime = time.Now()
	err = indexTaskRecordInfo.Update([]string{"status", "remark", "modify_time"})
	if err != nil {
		fmt.Println("修改子任务状态失败!")
		return
	}

	// 处理完成后标记任务状态
	defer func() {
		obj := models.IndexTaskRecord{}
		// 修改任务状态
		todoCount, tmpErr := obj.GetCountByCondition(fmt.Sprintf(` AND %s = ? AND %s = ? `, models.IndexTaskRecordColumns.IndexTaskID, models.IndexTaskRecordColumns.Status), []interface{}{indexTaskRecordInfo.IndexTaskID, `待处理`})
		if tmpErr != nil {
			err = fmt.Errorf("查找剩余任务数量失败, err: %s", tmpErr.Error())
			return
		}
		if todoCount <= 0 {
			indexTaskObj := models.IndexTask{}
			indexTaskInfo, tmpErr := indexTaskObj.GetByID(indexTaskRecordInfo.IndexTaskID)
			if tmpErr != nil {
				err = fmt.Errorf("查找任务失败, err: %s", tmpErr.Error())
				return
			}
			tmpUpdateCols := []string{`end_time`, "status", "update_time"}
			indexTaskInfo.EndTime = time.Now()
			indexTaskInfo.Status = `处理成功`
			indexTaskInfo.UpdateTime = time.Now()

			if indexTaskInfo.StartTime.IsZero() {
				indexTaskInfo.StartTime = time.Now()
				tmpUpdateCols = append(tmpUpdateCols, "start_time")
			}

			tmpErr = indexTaskInfo.Update(tmpUpdateCols)
			if tmpErr != nil {
				utils.FileLog.Error("标记任务状态失败, err: %s", tmpErr.Error())
			}
		}

		return
	}()

	// 修改模型状态
	switch taskType {
	case utils.INDEX_TASK_TYPE_AI_MODEL_TRAIN: // 训练模型
		// 修改模型状态信息
		if indexItem != nil {
			indexItem.TrainStatus = ai_predict_model.TrainStatusFailed
			indexItem.ModifyTime = time.Now()
			tmpErr := indexItem.Update([]string{"train_status", "modify_time"})
			if tmpErr != nil {
				utils.FileLog.Error("%d，修改模型训练状态失败, err: %s", indexItem.AiPredictModelIndexId, tmpErr.Error())
			}
		}

		// 修改模型配置状态信息
		if indexConfigItem != nil {
			indexConfigItem.TrainStatus = ai_predict_model.TrainStatusFailed
			indexConfigItem.Remark = errMsg
			indexConfigItem.ModifyTime = time.Now()
			tmpErr := indexConfigItem.Update([]string{"train_status", `remark`, "modify_time"})
			if tmpErr != nil {
				utils.FileLog.Error("%d，修改模型训练状态失败, err: %s", indexItem.AiPredictModelIndexId, tmpErr.Error())
			}
		}

	case utils.INDEX_TASK_TYPE_AI_MODEL_RUN: // 运行模型
		if indexItem != nil {
			indexItem.RunStatus = ai_predict_model.RunStatusFailed
			indexItem.ModifyTime = time.Now()
			tmpErr := indexItem.Update([]string{"run_status", "modify_time"})
			if tmpErr != nil {
				utils.FileLog.Error("%d，修改模型运行状态失败, err: %s", indexItem.AiPredictModelIndexId, tmpErr.Error())
			}
		}

	default:

		return
	}

	return
}

// HandleTaskRecordProcessingByTaskRecord
// @Description: 任务标记处理中
// @author: Roc
// @datetime 2025-05-09 16:24:38
// @param taskType string
// @param indexTaskRecordInfo *models.IndexTaskRecord
// @param indexConfigItem *ai_predict_model.AiPredictModelIndexConfig
// @param indexItem *ai_predict_model.AiPredictModelIndex
func HandleTaskRecordProcessingByTaskRecord(taskType string, indexTaskRecordInfo *models.IndexTaskRecord, indexConfigItem *ai_predict_model.AiPredictModelIndexConfig, indexItem *ai_predict_model.AiPredictModelIndex) {
	var err error
	defer func() {
		if err != nil {
			utils.FileLog.Error(fmt.Sprintf(`HandleTaskRecordFailByTaskRecord err:%v`, err))
		}
	}()

	// 修改子任务状态
	indexTaskRecordInfo.Status = `处理中`
	indexTaskRecordInfo.ModifyTime = time.Now()
	err = indexTaskRecordInfo.Update([]string{"status", "modify_time"})
	if err != nil {
		fmt.Println("修改子任务状态失败!")
		return
	}

	// 处理完成后标记任务状态
	defer func() {
		obj := models.IndexTaskRecord{}
		// 修改任务状态
		todoCount, tmpErr := obj.GetCountByCondition(fmt.Sprintf(` AND %s = ? AND %s = ? `, models.IndexTaskRecordColumns.IndexTaskID, models.IndexTaskRecordColumns.Status), []interface{}{indexTaskRecordInfo.IndexTaskID, `待处理`})
		if tmpErr != nil {
			err = fmt.Errorf("查找剩余任务数量失败, err: %s", tmpErr.Error())
			return
		}
		if todoCount <= 0 {
			indexTaskObj := models.IndexTask{}
			indexTaskInfo, tmpErr := indexTaskObj.GetByID(indexTaskRecordInfo.IndexTaskID)
			if tmpErr != nil {
				err = fmt.Errorf("查找任务失败, err: %s", tmpErr.Error())
				return
			}
			tmpUpdateCols := []string{`end_time`, "status", "update_time"}
			indexTaskInfo.EndTime = time.Now()
			indexTaskInfo.Status = `处理中`
			indexTaskInfo.UpdateTime = time.Now()

			if indexTaskInfo.StartTime.IsZero() {
				indexTaskInfo.StartTime = time.Now()
				tmpUpdateCols = append(tmpUpdateCols, "start_time")
			}

			tmpErr = indexTaskInfo.Update(tmpUpdateCols)
			if tmpErr != nil {
				utils.FileLog.Error("标记任务状态失败, err: %s", tmpErr.Error())
			}
		}

		return
	}()

	// 修改模型状态
	switch taskType {
	case utils.INDEX_TASK_TYPE_AI_MODEL_TRAIN: // 训练模型
		// 修改模型状态信息
		if indexItem != nil {
			indexItem.TrainStatus = ai_predict_model.TrainStatusTraining
			indexItem.ModifyTime = time.Now()
			tmpErr := indexItem.Update([]string{"train_status", "modify_time"})
			if tmpErr != nil {
				utils.FileLog.Error("%d，修改模型训练状态失败, err: %s", indexItem.AiPredictModelIndexId, tmpErr.Error())
			}
		}

		// 修改模型配置状态信息
		if indexConfigItem != nil {
			indexConfigItem.TrainStatus = ai_predict_model.TrainStatusTraining
			indexConfigItem.ModifyTime = time.Now()
			tmpErr := indexConfigItem.Update([]string{"train_status", "modify_time"})
			if tmpErr != nil {
				utils.FileLog.Error("%d，修改模型训练状态失败, err: %s", indexItem.AiPredictModelIndexId, tmpErr.Error())
			}
		}

	case utils.INDEX_TASK_TYPE_AI_MODEL_RUN: // 运行模型
		// 修改模型状态信息
		if indexItem != nil {
			indexItem.RunStatus = ai_predict_model.RunStatusRunning
			indexItem.ModifyTime = time.Now()
			tmpErr := indexItem.Update([]string{"run_status", "modify_time"})
			if tmpErr != nil {
				utils.FileLog.Error("%d，修改模型运行状态失败, err: %s", indexItem.AiPredictModelIndexId, tmpErr.Error())
			}
		}

	default:

		return
	}

	return
}