aiforge/routers/repo/modelarts.go

package repo

import (
	"code.gitea.io/gitea/modules/modelarts"
	"errors"
	"github.com/unknwon/com"
	"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/context"
	"code.gitea.io/gitea/modules/log"
	"code.gitea.io/gitea/modules/setting"
)

const (
	tplModelArtsIndex base.TplName = "repo/modelarts/index"
	tplModelArtsNew   base.TplName = "repo/modelarts/new"
	tplModelArtsShow  base.TplName = "repo/modelarts/show"
)

// MustEnableDataset check if repository enable internal cb
func MustEnableModelArts(ctx *context.Context) {
	if !ctx.Repo.CanRead(models.UnitTypeCloudBrain) {
		ctx.NotFound("MustEnableCloudbrain", nil)
		return
	}
}
func ModelArtsIndex(ctx *context.Context) {
	MustEnableModelArts(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.TypeCloudBrainTwo,
	})
	if err != nil {
		ctx.ServerError("Cloudbrain", err)
		return
	}

	for i, task := range ciTasks {
		if task.Status == string(models.JobRunning) {
			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, tplModelArtsIndex)
}

func ModelArtsNew(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

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

	ctx.Data["attachments"] = attachs
	ctx.Data["dataset_path"] = modelarts.DataSetMountPath
	ctx.Data["env"] = modelarts.NotebookEnv
	ctx.Data["notebook_type"] = modelarts.NotebookType
	ctx.Data["flavor"] = modelarts.FlavorInfo
	ctx.HTML(200, tplModelArtsNew)
}

func ModelArtsCreate(ctx *context.Context, form auth.CreateModelArtsForm) {
	ctx.Data["PageIsCloudBrain"] = true
	jobName := form.JobName
	uuid := form.Attachment
	description := form.Description
	//repo := ctx.Repo.Repository

	err := modelarts.GenerateTask(ctx, jobName, uuid, description)
	if err != nil {
		ctx.RenderWithErr(err.Error(), tplModelArtsNew, &form)
		return
	}

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

func ModelArtsShow(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()
		ctx.RenderWithErr(err.Error(), tplModelArtsIndex, nil)
		return
	}

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

	if result != nil {
		task.Status = result.Status
		err = models.UpdateJob(task)
		if err != nil {
			ctx.Data["error"] = err.Error()
			ctx.RenderWithErr(err.Error(), tplModelArtsIndex, nil)
			return
		}

		createTime, _ := com.StrTo(result.CreationTimestamp).Int64()
		result.CreateTime = time.Unix(int64(createTime/1000), 0).Format("2006-01-02 15:04:05")
		endTime, _ := com.StrTo(result.LatestUpdateTimestamp).Int64()
		result.LatestUpdateTime = time.Unix(int64(endTime/1000), 0).Format("2006-01-02 15:04:05")
		result.QueuingInfo.BeginTime = time.Unix(int64(result.QueuingInfo.BeginTimestamp/1000), 0).Format("2006-01-02 15:04:05")
		result.QueuingInfo.EndTime = time.Unix(int64(result.QueuingInfo.EndTimestamp/1000), 0).Format("2006-01-02 15:04:05")
	}

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

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

	result, err := modelarts.GetJob(jobID)
	if err != nil {
		ctx.RenderWithErr(err.Error(), tplModelArtsIndex, nil)
		return
	}

	res, err := modelarts.GetJobToken(jobID)
	if err != nil {
		ctx.RenderWithErr(err.Error(), tplModelArtsIndex, nil)
		return
	}

	urls := strings.Split(result.Spec.Annotations.Url, "/")
	urlPrefix := result.Spec.Annotations.TargetDomain
	for i, url := range urls {
		if i > 2 {
			urlPrefix += "/" + url
		}
	}

	//urlPrefix := result.Spec.Annotations.TargetDomain + "/modelarts/internal/hub/notebook/user/" + task.JobID
	log.Info(urlPrefix)
	debugUrl := urlPrefix + "?token=" + res.Token
	ctx.Redirect(debugUrl)
}

func ModelArtsStop(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
	}

	param := models.NotebookAction{
		Action: models.ActionStop,
	}
	res, err := modelarts.StopJob(jobID, param)
	if err != nil {
		log.Error("StopJob(%s) failed:%v", task.JobName, err.Error())
		ctx.ServerError("StopJob failed", err)
		return
	}

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

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

func ModelArtsDel(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 = modelarts.DelJob(jobID)
	if err != nil {
		log.Error("DelJob(%s) failed:%v", task.JobName, err.Error())
		ctx.ServerError("DelJob failed", err)
		return
	}

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

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