aiforge/services/cloudbrain/cloudbrainTask/sync_status.go

package cloudbrainTask

import (
	"net/http"
	"os"
	"strings"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/cloudbrain"
	"code.gitea.io/gitea/modules/httplib"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/notification"
	"code.gitea.io/gitea/modules/setting"
	"code.gitea.io/gitea/modules/storage"
)

var noteBookOKMap = make(map[int64]int, 20)

//if a task notebook url can get  two times,  the notebook can browser.
const successfulCount = 3

func SyncCloudBrainOneStatus(task *models.Cloudbrain) (*models.Cloudbrain, error) {
	jobResult, err := cloudbrain.GetJob(task.JobID)
	if err != nil {

		log.Error("GetJob failed:", err)

		return task, err
	}
	result, err := models.ConvertToJobResultPayload(jobResult.Payload)
	if err != nil {
		log.Error("ConvertToJobResultPayload failed:", err)
		return task, err
	}
	oldStatus := task.Status

	if result.JobStatus.State != string(models.JobWaiting) && result.JobStatus.State != string(models.JobFailed) {
		taskRoles := result.TaskRoles
		taskRes, _ := models.ConvertToTaskPod(taskRoles[cloudbrain.SubTaskName].(map[string]interface{}))

		task.ContainerIp = taskRes.TaskStatuses[0].ContainerIP
		task.ContainerID = taskRes.TaskStatuses[0].ContainerID
	}

	if (result.JobStatus.State != string(models.JobWaiting) && result.JobStatus.State != string(models.JobRunning)) ||
		task.Status == string(models.JobRunning) || (result.JobStatus.State == string(models.JobRunning) && isNoteBookReady(task)) {

		models.ParseAndSetDurationFromCloudBrainOne(result, task)
		task.Status = result.JobStatus.State
		if oldStatus != task.Status {
			notification.NotifyChangeCloudbrainStatus(task, oldStatus)
			err := updateLogFile(task, result)
			if err != nil {
				log.Error("updateLogFile failed:", err)
				return task, err
			}
		}
		err = models.UpdateJob(task)
		if err != nil {
			log.Error("UpdateJob failed:", err)
			return task, err
		}
	}
	return task, nil

}

func updateLogFile(task *models.Cloudbrain, result models.JobResultPayload) error {
	if task.Type == models.TypeCloudBrainOne {
		taskRoles := result.TaskRoles
		taskRes, _ := models.ConvertToTaskPod(taskRoles[cloudbrain.SubTaskName].(map[string]interface{}))
		existStr := taskRes.TaskStatuses[0].ExitDiagnostics
		logDir := "/model"
		files, err := storage.GetOneLevelAllObjectUnderDirMinio(setting.Attachment.Minio.Bucket, setting.CBCodePathPrefix+task.JobName+logDir, "")
		if err != nil {
			log.Error("query cloudbrain model failed: %v", err)
			return err
		}
		fileName := ""
		for _, file := range files {
			if strings.HasSuffix(file.FileName, "log.txt") {
				fileName = file.FileName
				break
			}
		}
		if fileName != "" {
			prefix := "/" + setting.CBCodePathPrefix + task.JobName + "/model"
			configFile, err := os.OpenFile(setting.Attachment.Minio.RealPath+setting.Attachment.Minio.Bucket+prefix+"/"+fileName, os.O_WRONLY|os.O_APPEND, 0666)
			if err != nil {
				log.Error("open file(%s) failed:%s", prefix+fileName, err)
				return err
			}
			defer configFile.Close()
			_, err = configFile.WriteString(existStr)
			if err != nil {
				log.Error("WriteString failed:%v", err)
				return err
			}
		}

	}
	return nil
}

func isNoteBookReady(task *models.Cloudbrain) bool {
	if task.JobType != string(models.JobTypeDebug) {
		return true
	}
	noteBookUrl := setting.DebugServerHost + "jpylab_" + task.JobID + "_" + task.SubTaskName
	r := httplib.Get(noteBookUrl)
	res, err := r.Response()
	if err != nil {
		return false
	}
	if res.StatusCode == http.StatusOK {
		count := noteBookOKMap[task.ID]
		if count < successfulCount-1 {
			noteBookOKMap[task.ID] = count + 1
			return false
		} else {
			delete(noteBookOKMap, task.ID)
			return true
		}

	}
	return false

}