aiforge/modules/cloudbrain/cloudbrain.go

package cloudbrain

import (
	"errors"

	"code.gitea.io/gitea/modules/setting"

	"code.gitea.io/gitea/models"
	"code.gitea.io/gitea/modules/context"
	"code.gitea.io/gitea/modules/log"
)

const (
	Command               = `pip3 install jupyterlab==2.2.5 -i https://pypi.tuna.tsinghua.edu.cn/simple;service ssh stop;jupyter lab --no-browser --ip=0.0.0.0 --allow-root --notebook-dir="/code" --port=80 --LabApp.token="" --LabApp.allow_origin="self https://cloudbrain.pcl.ac.cn"`
	CodeMountPath         = "/code"
	DataSetMountPath      = "/dataset"
	ModelMountPath        = "/model"
	BenchMarkMountPath    = "/benchmark"
	Snn4imagenetMountPath = "/snn4imagenet"
	BrainScoreMountPath   = "/brainscore"
	TaskInfoName          = "/taskInfo"

	SubTaskName = "task1"

	Success = "S000"
)

var (
	ResourceSpecs *models.ResourceSpecs
)

func isAdminOrOwnerOrJobCreater(ctx *context.Context, job *models.Cloudbrain, err error) bool {

	if err != nil {
		return ctx.IsUserRepoOwner() || ctx.IsUserSiteAdmin()
	} else {
		return ctx.IsUserRepoOwner() || ctx.IsUserSiteAdmin() || ctx.User.ID == job.UserID
	}

}

func CanDeleteDebugJob(ctx *context.Context, job *models.Cloudbrain) bool {

	return isAdminOrOwnerOrJobCreater(ctx, job, nil)
}

func CanDeleteTrainJob(ctx *context.Context, job *models.Cloudbrain) bool {

	return isAdminOrOwnerOrJobCreater(ctx, job, nil)
}

func CanCreateOrDebugJob(ctx *context.Context) bool {
	return ctx.Repo.CanWrite(models.UnitTypeCloudBrain)
}

func CanModifyJob(ctx *context.Context, job *models.Cloudbrain) bool {

	return isAdminOrJobCreater(ctx, job, nil)
}

func isAdminOrJobCreater(ctx *context.Context, job *models.Cloudbrain, err error) bool {

	if err != nil {
		return ctx.IsUserSiteAdmin()
	} else {
		return ctx.IsUserSiteAdmin() || ctx.User.ID == job.UserID
	}

}

func AdminOrOwnerOrJobCreaterRight(ctx *context.Context) {

	var jobID = ctx.Params(":jobid")

	job, err := models.GetCloudbrainByJobID(jobID)

	if !isAdminOrOwnerOrJobCreater(ctx, job, err) {

		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
	}

}

func AdminOrJobCreaterRight(ctx *context.Context) {

	var jobID = ctx.Params(":jobid")
	job, err := models.GetCloudbrainByJobID(jobID)
	if !isAdminOrJobCreater(ctx, job, err) {

		ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
	}

}

func GenerateTask(ctx *context.Context, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, snn4imagenetPath, brainScorePath, jobType, gpuQueue string, resourceSpecId int) error {
	dataActualPath := setting.Attachment.Minio.RealPath +
		setting.Attachment.Minio.Bucket + "/" +
		setting.Attachment.Minio.BasePath +
		models.AttachmentRelativePath(uuid) +
		uuid

	var resourceSpec *models.ResourceSpec
	for _, spec := range ResourceSpecs.ResourceSpec {
		if resourceSpecId == spec.Id {
			resourceSpec = spec
		}
	}

	if resourceSpec == nil {
		log.Error("no such resourceSpecId(%d)", resourceSpecId, ctx.Data["MsgID"])
		return errors.New("no such resourceSpec")
	}

	jobResult, err := CreateJob(jobName, models.CreateJobParams{
		JobName:    jobName,
		RetryCount: 1,
		GpuType:    gpuQueue,
		Image:      image,
		TaskRoles: []models.TaskRole{
			{
				Name:                  SubTaskName,
				TaskNumber:            1,
				MinSucceededTaskCount: 1,
				MinFailedTaskCount:    1,
				CPUNumber:             resourceSpec.CpuNum,
				GPUNumber:             resourceSpec.GpuNum,
				MemoryMB:              resourceSpec.MemMiB,
				ShmMB:                 resourceSpec.ShareMemMiB,
				Command:               command,
				NeedIBDevice:          false,
				IsMainRole:            false,
				UseNNI:                false,
			},
		},
		Volumes: []models.Volume{
			{
				HostPath: models.StHostPath{
					Path:      codePath,
					MountPath: CodeMountPath,
					ReadOnly:  false,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:      dataActualPath,
					MountPath: DataSetMountPath,
					ReadOnly:  true,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:      modelPath,
					MountPath: ModelMountPath,
					ReadOnly:  false,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:      benchmarkPath,
					MountPath: BenchMarkMountPath,
					ReadOnly:  true,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:      snn4imagenetPath,
					MountPath: Snn4imagenetMountPath,
					ReadOnly:  true,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:      brainScorePath,
					MountPath: BrainScoreMountPath,
					ReadOnly:  true,
				},
			},
		},
	})
	if err != nil {
		log.Error("CreateJob failed:", err.Error())
		return err
	}
	if jobResult.Code != Success {
		log.Error("CreateJob(%s) failed:%s", jobName, jobResult.Msg)
		return errors.New(jobResult.Msg)
	}

	var jobID = jobResult.Payload["jobId"].(string)
	err = models.CreateCloudbrain(&models.Cloudbrain{
		Status:      string(models.JobWaiting),
		UserID:      ctx.User.ID,
		RepoID:      ctx.Repo.Repository.ID,
		JobID:       jobID,
		JobName:     jobName,
		SubTaskName: SubTaskName,
		JobType:     jobType,
		Type:        models.TypeCloudBrainOne,
		Uuid:        uuid,
	})

	if err != nil {
		return err
	}

	return nil
}