aiforge/routers/repo/cloudbrain.go

package repo

import (
	"code.gitea.io/gitea/modules/git"
	"encoding/json"
	"errors"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/auth"
	"code.gitea.io/gitea/modules/base"
	"code.gitea.io/gitea/modules/cloudbrain"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
)

const (
	tplCloudBrainIndex base.TplName = "repo/cloudbrain/index"
	tplCloudBrainNew   base.TplName = "repo/cloudbrain/new"
	tplCloudBrainShow  base.TplName = "repo/cloudbrain/show"
)

// MustEnableDataset check if repository enable internal cb
func MustEnableCloudbrain(ctx *context.Context) {
	if !ctx.Repo.CanRead(models.UnitTypeCloudBrain) {
		ctx.NotFound("MustEnableCloudbrain", nil)
		return
	}
}
func CloudBrainIndex(ctx *context.Context) {
	MustEnableCloudbrain(ctx)
	repo := ctx.Repo.Repository
	page := ctx.QueryInt("page")
	if page <= 0 {
		page = 1
	}

	ciTasks, count, err := models.Cloudbrains(&models.CloudbrainsOptions{
		ListOptions: models.ListOptions{
			Page:     page,
			PageSize: setting.UI.IssuePagingNum,
		},
		RepoID: repo.ID,
		Type:	models.TypeCloudBrainOne,
	})
	if err != nil {
		ctx.ServerError("Cloudbrain", err)
		return
	}

	timestamp := time.Now().Unix()
	for i, task := range ciTasks {
		if task.Status == string(models.JobRunning) && (timestamp-int64(task.CreatedUnix) > 30) {
			ciTasks[i].CanDebug = true
		} else {
			ciTasks[i].CanDebug = false
		}
	}

	pager := context.NewPagination(int(count), setting.UI.IssuePagingNum, page, 5)
	pager.SetDefaultParams(ctx)
	ctx.Data["Page"] = pager

	ctx.Data["PageIsCloudBrain"] = true
	ctx.Data["Tasks"] = ciTasks
	ctx.HTML(200, tplCloudBrainIndex)
}

func cutString(str string, lens int) string {
	if len(str) < lens {
		return str
	}
	return str[:lens]
}

func CloudBrainNew(ctx *context.Context) {
	ctx.Data["PageIsCloudBrain"] = true

	t := time.Now()
	var jobName = cutString(ctx.User.Name, 5) + t.Format("2006010215") + strconv.Itoa(int(t.Unix()))[5:]
	ctx.Data["job_name"] = jobName

	result, err := cloudbrain.GetImages()
	if err != nil {
		ctx.Data["error"] = err.Error()
	}

	for i, payload := range result.Payload {
		if strings.HasPrefix(result.Payload[i].Place, "192.168") {
			result.Payload[i].PlaceView = payload.Place[strings.Index(payload.Place, "/"):len(payload.Place)]
		} else {
			result.Payload[i].PlaceView = payload.Place
		}
	}

	ctx.Data["images"] = result.Payload

	attachs, err := models.GetAllUserAttachments(ctx.User.ID)
	if err != nil {
		ctx.ServerError("GetAllUserAttachments failed:", err)
		return
	}

	ctx.Data["attachments"] = attachs
	ctx.Data["command"] = cloudbrain.Command
	ctx.Data["code_path"] = cloudbrain.CodeMountPath
	ctx.Data["dataset_path"] = cloudbrain.DataSetMountPath
	ctx.Data["model_path"] = cloudbrain.ModelMountPath
	ctx.Data["benchmark_path"] = cloudbrain.BenchMarkMountPath
	ctx.Data["is_benchmark_enabled"] = setting.IsBenchmarkEnabled
	ctx.HTML(200, tplCloudBrainNew)
}

func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
	ctx.Data["PageIsCloudBrain"] = true
	jobName := form.JobName
	image := form.Image
	command := form.Command
	uuid := form.Attachment
	jobType := form.JobType
	codePath := setting.JobPath + jobName + cloudbrain.CodeMountPath

	if jobType != string(models.JobTypeBenchmark) && jobType != string(models.JobTypeDebug) {
		log.Error("jobtype error:", jobType)
		ctx.RenderWithErr("jobtype error", tplCloudBrainNew, &form)
		return
	}
	repo := ctx.Repo.Repository
	downloadCode(repo, codePath)

	modelPath := setting.JobPath + jobName + cloudbrain.ModelMountPath
	err := os.MkdirAll(modelPath, os.ModePerm)
	if err != nil {
		ctx.RenderWithErr(err.Error(), tplCloudBrainNew, &form)
		return
	}

	benchmarkPath := setting.JobPath + jobName + cloudbrain.BenchMarkMountPath
	if setting.IsBenchmarkEnabled && jobType == string(models.JobTypeBenchmark) {
		downloadBenchmarkCode(repo, jobName, benchmarkPath)
	}

	err = cloudbrain.GenerateTask(ctx, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, jobType)
	if err != nil {
		ctx.RenderWithErr(err.Error(), tplCloudBrainNew, &form)
		return
	}
	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain")
}

func CloudBrainShow(ctx *context.Context) {
	ctx.Data["PageIsCloudBrain"] = true

	var jobID = ctx.Params(":jobid")
	task, err := models.GetCloudbrainByJobID(jobID)
	if err != nil {
		ctx.Data["error"] = err.Error()
	}

	result, err := cloudbrain.GetJob(jobID)
	if err != nil {
		ctx.Data["error"] = err.Error()
	}

	if result != nil {
		jobRes, _ := models.ConvertToJobResultPayload(result.Payload)
		ctx.Data["result"] = jobRes
		taskRoles := jobRes.TaskRoles
		taskRes, _ := models.ConvertToTaskPod(taskRoles[cloudbrain.SubTaskName].(map[string]interface{}))
		ctx.Data["taskRes"] = taskRes
		task.Status = taskRes.TaskStatuses[0].State
		task.ContainerID = taskRes.TaskStatuses[0].ContainerID
		task.ContainerIp = taskRes.TaskStatuses[0].ContainerIP
		err = models.UpdateJob(task)
		if err != nil {
			ctx.Data["error"] = err.Error()
		}
	}

	ctx.Data["task"] = task
	ctx.Data["jobID"] = jobID
	ctx.HTML(200, tplCloudBrainShow)
}

func CloudBrainDebug(ctx *context.Context) {
	var jobID = ctx.Params(":jobid")
	task, err := models.GetCloudbrainByJobID(jobID)
	if err != nil {
		ctx.ServerError("GetCloudbrainByJobID failed", err)
		return
	}

	debugUrl := setting.DebugServerHost + "jpylab_" + task.JobID + "_" + task.SubTaskName
	ctx.Redirect(debugUrl)
}

func CloudBrainCommitImage(ctx *context.Context, form auth.CommitImageCloudBrainForm) {
	var jobID = ctx.Params(":jobid")
	task, err := models.GetCloudbrainByJobID(jobID)
	if err != nil {
		ctx.JSON(200, map[string]string{
			"result_code": "-1",
			"error_msg":   "GetCloudbrainByJobID failed",
		})
		return
	}

	err = cloudbrain.CommitImage(jobID, models.CommitImageParams{
		Ip:               task.ContainerIp,
		TaskContainerId:  task.ContainerID,
		ImageDescription: form.Description,
		ImageTag:         form.Tag,
	})
	if err != nil {
		log.Error("CommitImage(%s) failed:", task.JobName, err.Error())
		ctx.JSON(200, map[string]string{
			"result_code": "-1",
			"error_msg":   "CommitImage failed",
		})
		return
	}

	ctx.JSON(200, map[string]string{
		"result_code": "0",
		"error_msg":   "",
	})
}

func CloudBrainStop(ctx *context.Context) {
	var jobID = ctx.Params(":jobid")
	log.Info(jobID)
	task, err := models.GetCloudbrainByJobID(jobID)
	if err != nil {
		ctx.ServerError("GetCloudbrainByJobID failed", err)
		return
	}

	if task.Status != string(models.JobRunning) {
		log.Error("the job(%s) is not running", task.JobName)
		ctx.ServerError("the job is not running", errors.New("the job is not running"))
		return
	}

	err = cloudbrain.StopJob(jobID)
	if err != nil {
		log.Error("StopJob(%s) failed:%v", task.JobName, err.Error())
		ctx.ServerError("StopJob failed", err)
		return
	}

	task.Status = string(models.JobStopped)
	err = models.UpdateJob(task)
	if err != nil {
		ctx.ServerError("UpdateJob failed", err)
		return
	}

	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain")
}

func CloudBrainDel(ctx *context.Context) {
	var jobID = ctx.Params(":jobid")
	task, err := models.GetCloudbrainByJobID(jobID)
	if err != nil {
		ctx.ServerError("GetCloudbrainByJobID failed", err)
		return
	}

	if task.Status != string(models.JobStopped) {
		log.Error("the job(%s) has not been stopped", task.JobName)
		ctx.ServerError("the job has not been stopped", errors.New("the job has not been stopped"))
		return
	}

	err = models.DeleteJob(task)
	if err != nil {
		ctx.ServerError("DeleteJob failed", err)
		return
	}

	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/cloudbrain")
}

func CloudBrainBenchmark(ctx *context.Context) {
	var jobID = ctx.Params(":jobid")
	_, err := models.GetCloudbrainByJobID(jobID)
	if err != nil {
		ctx.ServerError("GetCloudbrainByJobID failed", err)
		return
	}

	ctx.Redirect(setting.BenchmarkServerHost)
}

func downloadCode(repo *models.Repository, codePath string) error {
	if err := git.Clone(repo.RepoPath(), codePath, git.CloneRepoOptions{}); err != nil {
		log.Error("Failed to clone repository: %s (%v)", repo.FullName(), err)
		return err
	}

	return nil
}

func downloadBenchmarkCode(repo *models.Repository, taskName string, benchmarkPath string) error {
	err := os.MkdirAll(benchmarkPath, os.ModePerm)
	if err != nil {
		log.Error("mkdir benchmarkPath failed", err.Error())
		return err
	}

	command := "git clone " + setting.BenchmarkCode + " " + benchmarkPath
	cmd := exec.Command("/bin/bash", "-c", command)
	output, err := cmd.Output()
	log.Info(string(output))
	if err != nil {
		log.Error("exec.Command(%s) failed:%v", command, err)
		return err
	}

	fileName := benchmarkPath + cloudbrain.TaskInfoName
	f, err := os.OpenFile(fileName, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, os.ModePerm)
	if err != nil {
		log.Error("OpenFile failed", err.Error())
		return err
	}

	defer f.Close()

	data, err := json.Marshal(models.TaskInfo{
		Username: repo.Owner.Name,
		TaskName: taskName,
		CodeName: repo.Name,
	})
	if err != nil {
		log.Error("json.Marshal failed", err.Error())
		return err
	}

	_, err = f.Write(data)
	if err != nil {
		log.Error("WriteString failed", err.Error())
		return err
	}

	return nil
}