aiforge/services/cloudbrain/cloudbrainTask/sync_status.go

package cloudbrainTask

import (
	"net/http"

	"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"
)

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

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

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 = models.UpdateJob(task)
		if err != nil {
			log.Error("UpdateJob failed:", err)
			return task, err
		}
	}
	return task, 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 || (noteBookFailMap[task.ID]==0 || count < maxSuccessfulCount-1){
			noteBookOKMap[task.ID] = count + 1
			return false
		} else {
			delete(noteBookOKMap, task.ID)
			delete(noteBookFailMap, task.ID)
			return true
		}

	}else{
		noteBookFailMap[task.ID]+=1
		if noteBookOKMap[task.ID]>0{
			noteBookOKMap[task.ID]=0
		}
	}
	return false

}