aiforge/routers/api/v1/repo/cloudbrain.go

// Copyright 2016 The Gogs Authors. All rights reserved.
// Copyright 2018 The Gitea Authors. All rights reserved.
// Use of this source code is governed by a MIT-style
// license that can be found in the LICENSE file.

package repo

import (
	"encoding/json"
	"net/http"
	"sort"
	"strings"
	"time"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/cloudbrain"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/storage"
	routerRepo "code.gitea.io/gitea/routers/repo"
)

// cloudbrain get job task by jobid
func GetCloudbrainTask(ctx *context.APIContext) {
	// swagger:operation GET /repos/{owner}/{repo}/cloudbrain/{jobid} cloudbrain jobTask
	// ---
	// summary: Get a single task
	// produces:
	// - application/json
	// parameters:
	// - name: owner
	//   in: path
	//   description: owner of the repo
	//   type: string
	//   required: true
	// - name: repo
	//   in: path
	//   description: name of the repo
	//   type: string
	//   required: true
	// - name: jobid
	//   in: path
	//   description: id of cloudbrain jobid
	//   type: string
	//   required: true
	// responses:
	//   "200":

	//     "$ref": "#/responses/Label"

	var (
		err error
	)

	ID := ctx.Params(":id")
	job, err := models.GetCloudbrainByID(ID)
	if err != nil {
		ctx.NotFound(err)
		log.Error("GetCloudbrainByID failed:", err)
		return
	}
	jobResult, err := cloudbrain.GetJob(job.JobID)
	if err != nil {
		ctx.NotFound(err)
		log.Error("GetJob failed:", err)
		return
	}
	result, _ := models.ConvertToJobResultPayload(jobResult.Payload)
	if err != nil {
		ctx.NotFound(err)
		log.Error("ConvertToJobResultPayload failed:", err)
		return
	}

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

	if result.JobStatus.State != string(models.JobWaiting) {
		models.ParseAndSetDurationFromCloudBrainOne(result, job)
		err = models.UpdateJob(job)
		if err != nil {
			log.Error("UpdateJob failed:", err)
		}
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"ID":            ID,
		"JobName":       result.Config.JobName,
		"JobStatus":     result.JobStatus.State,
		"SubState":      result.JobStatus.SubState,
		"CreatedTime":   time.Unix(result.JobStatus.CreatedTime/1000, 0).Format("2006-01-02 15:04:05"),
		"CompletedTime": time.Unix(result.JobStatus.CompletedTime/1000, 0).Format("2006-01-02 15:04:05"),
	})

}

func GetCloudBrainInferenceJob(ctx *context.APIContext) {

	jobID := ctx.Params(":jobid")
	job, err := models.GetCloudbrainByJobID(jobID)
	if err != nil {
		ctx.NotFound(err)
		return
	}
	jobResult, err := cloudbrain.GetJob(job.JobID)
	if err != nil {
		ctx.NotFound(err)
		log.Error("GetJob failed:", err)
		return
	}
	result, err := models.ConvertToJobResultPayload(jobResult.Payload)
	if err != nil {
		ctx.NotFound(err)
		log.Error("ConvertToJobResultPayload failed:", err)
		return
	}

	job.Status = result.JobStatus.State
	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{}))

		job.ContainerIp = taskRes.TaskStatuses[0].ContainerIP
		job.ContainerID = taskRes.TaskStatuses[0].ContainerID
		job.Status = taskRes.TaskStatuses[0].State
	}

	if result.JobStatus.State != string(models.JobWaiting) {
		models.ParseAndSetDurationFromCloudBrainOne(result, job)
		err = models.UpdateJob(job)
		if err != nil {
			log.Error("UpdateJob failed:", err)
		}
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":       jobID,
		"JobStatus":   job.Status,
		"JobDuration": job.TrainJobDuration,
	})

}

func DelCloudBrainJob(ctx *context.APIContext) {
	jobID := ctx.Params(":jobid")

	errStr := cloudbrain.DelCloudBrainJob(jobID)

	if errStr != "" {
		ctx.JSON(http.StatusOK, models.BaseErrorMessage(ctx.Tr(errStr)))
	} else {
		ctx.JSON(http.StatusOK, models.BaseOKMessage)
	}

}

func InferencJobResultList(ctx *context.APIContext) {
	jobID := ctx.Params(":jobid")
	parentDir := ctx.Query("parentDir")
	dirArray := strings.Split(parentDir, "/")

	task, err := models.GetCloudbrainByJobID(jobID)
	if err != nil {
		log.Error("get cloud brain err:", err)
		ctx.ServerError("get cloud brain information failed:", err)

	}

	//get dirs
	dirs, err := routerRepo.GetResultDirs(task.JobName, parentDir)
	if err != nil {
		log.Error("GetModelDirs failed:%v", err.Error(), ctx.Data["msgID"])
		ctx.ServerError("GetModelDirs failed:", err)
		return
	}

	var fileInfos []storage.FileInfo
	err = json.Unmarshal([]byte(dirs), &fileInfos)
	if err != nil {
		log.Error("json.Unmarshal failed:%v", err.Error(), ctx.Data["msgID"])
		ctx.ServerError("json.Unmarshal failed:", err)
		return
	}

	for i, fileInfo := range fileInfos {
		temp, _ := time.Parse("2006-01-02 15:04:05", fileInfo.ModTime)
		fileInfos[i].ModTime = temp.Local().Format("2006-01-02 15:04:05")
	}

	sort.Slice(fileInfos, func(i, j int) bool {
		return fileInfos[i].ModTime > fileInfos[j].ModTime
	})

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":            jobID,
		"StatusOK":         0,
		"Path":             dirArray,
		"Dirs":             fileInfos,
		"task":             task,
		"PageIsCloudBrain": true,
	})

}

func CloudbrainGetLog(ctx *context.Context) {
	ID := ctx.Params(":id")
	job, err := models.GetCloudbrainByID(ID)
	if err != nil {
		log.Error("GetCloudbrainByJobName failed: %v", err, ctx.Data["MsgID"])
		ctx.ServerError(err.Error(), err)
		return
	}

	var hits []models.Hits
	result, err := cloudbrain.GetJobLog(job.JobID)
	if err != nil {
		log.Error("GetJobLog failed: %v", err, ctx.Data["MsgID"])
		ctx.ServerError(err.Error(), err)
		return
	}
	hits = result.Hits.Hits

	//if the size equal page_size, then take the scroll_id to get all log and delete the scroll_id(the num of scroll_id is limited)
	if len(result.Hits.Hits) >= cloudbrain.LogPageSize {
		for {
			resultNext, err := cloudbrain.GetJobAllLog(result.ScrollID)
			if err != nil {
				log.Error("GetJobAllLog failed: %v", err, ctx.Data["MsgID"])
			} else {
				for _, hit := range resultNext.Hits.Hits {
					hits = append(hits, hit)
				}
			}

			if len(resultNext.Hits.Hits) < cloudbrain.LogPageSize {
				log.Info("get all log already")
				break
			}
		}
	}

	cloudbrain.DeleteJobLogToken(result.ScrollID)

	sort.Slice(hits, func(i, j int) bool {
		return hits[i].Sort[0] < hits[j].Sort[0]
	})

	var content string
	for _, log := range hits {
		content += log.Source.Message + "\n"
	}

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobName": job.JobName,
		"Content": content,
	})

	return
}

func CloudBrainModelList(ctx *context.APIContext) {
	var (
		err error
	)

	var jobID = ctx.Params(":jobid")
	var versionName = ctx.Query("version_name")
	parentDir := ctx.Query("parentDir")
	dirArray := strings.Split(parentDir, "/")

	task, err := models.GetCloudbrainByJobIDAndVersionName(jobID, versionName)
	if err != nil {
		log.Error("GetCloudbrainByJobID(%s) failed:%v", task.JobName, err.Error())
		return
	}

	//get dirs
	dirs, err := routerRepo.GetModelDirs(task.JobName, parentDir)
	if err != nil {
		log.Error("GetModelDirs failed:%v", err.Error(), ctx.Data["msgID"])
		ctx.ServerError("GetModelDirs failed:", err)
		return
	}

	var fileInfos []storage.FileInfo
	err = json.Unmarshal([]byte(dirs), &fileInfos)
	if err != nil {
		log.Error("json.Unmarshal failed:%v", err.Error(), ctx.Data["msgID"])
		ctx.ServerError("json.Unmarshal failed:", err)
		return
	}

	for i, fileInfo := range fileInfos {
		temp, _ := time.Parse("2006-01-02 15:04:05", fileInfo.ModTime)
		fileInfos[i].ModTime = temp.Local().Format("2006-01-02 15:04:05")
	}

	sort.Slice(fileInfos, func(i, j int) bool {
		return fileInfos[i].ModTime > fileInfos[j].ModTime
	})

	ctx.JSON(http.StatusOK, map[string]interface{}{
		"JobID":            jobID,
		"VersionName":      versionName,
		"StatusOK":         0,
		"Path":             dirArray,
		"Dirs":             fileInfos,
		"task":             task,
		"PageIsCloudBrain": true,
	})
}