aiforge/modules/cloudbrain/cloudbrain.go

package cloudbrain

import (
	"errors"

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

const (
	Command = `pip3 install jupyterlab==1.1.4;service ssh stop;jupyter lab --no-browser --ip=0.0.0.0 --allow-root --notebook-dir=\"/userhome\" --port=80  --NotebookApp.token=\"\" --LabApp.allow_origin=\"self https://cloudbrain.pcl.ac.cn\"`
	CodeMountPath =	"/code"
	DataSetMountPath = "/dataset"
	ModelMountPath = "/model"

	SubTaskName = "task1"

	DebugGPUType = "DEBUG"

)

func GenerateTask(ctx *context.Context, jobName, image, command string) error {
	jobResult, err := CreateJob(jobName, models.CreateJobParams{
		JobName:    jobName,
		RetryCount: 1,
		GpuType:    DebugGPUType,
		Image:      image,
		TaskRoles: []models.TaskRole{
			{
				Name:                  SubTaskName,
				TaskNumber:            1,
				MinSucceededTaskCount: 1,
				MinFailedTaskCount:    1,
				CPUNumber:             2,
				GPUNumber:             1,
				MemoryMB:              16384,
				ShmMB:                 8192,
				Command:               command,
				NeedIBDevice:          false,
				IsMainRole:            false,
				UseNNI:				   false,
			},
		},
		Volumes: []models.Volume{
			{
				HostPath: models.StHostPath{
					Path:				"",
					MountPath:			CodeMountPath,
					ReadOnly: 			true,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:				"",
					MountPath:			DataSetMountPath,
					ReadOnly: 			false,
				},
			},
			{
				HostPath: models.StHostPath{
					Path:				"",
					MountPath:			ModelMountPath,
					ReadOnly: 			true,
				},
			},
		},
	})
	if err != nil {
		return err
	}
	if jobResult.Code != "S000" {
		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,
	})

	if err != nil {
		return err
	}

	return nil
}