gpu model list

routers/api/v1/repo/modelarts.go

@@ -7,6 +7,7 @@ package repo
 
 import (
 	"code.gitea.io/gitea/modules/grampus"
+	"encoding/json"
 	"net/http"
 	"strconv"
 	"strings"
@@ -397,11 +398,29 @@ func ModelList(ctx *context.APIContext) {
 		log.Error("GetCloudbrainByJobID(%s) failed:%v", task.JobName, err.Error())
 		return
 	}
-	models, err := storage.GetObsListObject(task.JobName, "output/", parentDir, versionName)
-	if err != nil {
-		log.Info("get TrainJobListModel failed:", err)
-		ctx.ServerError("GetObsListObject:", err)
-		return
+
+	var fileInfos []storage.FileInfo
+	if task.ComputeResource == models.NPUResource {
+		fileInfos, err = storage.GetObsListObject(task.JobName, "output/", parentDir, versionName)
+		if err != nil {
+			log.Info("get TrainJobListModel failed:", err)
+			ctx.ServerError("GetObsListObject:", err)
+			return
+		}
+	} else if task.ComputeResource == models.GPUResource {
+		files, err := routerRepo.GetModelDirs(task.JobName, parentDir)
+		if err != nil {
+			log.Info("GetModelDirs failed:", err)
+			ctx.ServerError("GetModelDirs:", err)
+			return
+		}
+
+		err = json.Unmarshal([]byte(files), &fileInfos)
+		if err != nil {
+			log.Error("json.Unmarshal failed:%v", err.Error(), ctx.Data["msgID"])
+			ctx.ServerError("json.Unmarshal failed:", err)
+			return
+		}
 	}
 
 	ctx.JSON(http.StatusOK, map[string]interface{}{
@@ -409,7 +428,7 @@ func ModelList(ctx *context.APIContext) {
 		"VersionName":      versionName,
 		"StatusOK":         0,
 		"Path":             dirArray,
-		"Dirs":             models,
+		"Dirs":             fileInfos,
 		"task":             task,
 		"PageIsCloudBrain": true,
 	})

routers/repo/modelarts.go

@@ -2240,14 +2240,25 @@ func ModelDownload(ctx *context.Context) {
 		return
 	}
 
-	path := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, task.JobName, setting.OutPutPath, versionName, parentDir, fileName), "/")
-
-	url, err := storage.GetObsCreateSignedUrlByBucketAndKey(setting.Bucket, path)
-	if err != nil {
-		log.Error("GetObsCreateSignedUrl failed: %v", err.Error(), ctx.Data["msgID"])
-		ctx.ServerError("GetObsCreateSignedUrl", err)
-		return
+	var url string
+	if task.ComputeResource == models.NPUResource {
+		path := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, task.JobName, setting.OutPutPath, versionName, parentDir, fileName), "/")
+		url, err = storage.GetObsCreateSignedUrlByBucketAndKey(setting.Bucket, path)
+		if err != nil {
+			log.Error("GetObsCreateSignedUrl failed: %v", err.Error(), ctx.Data["msgID"])
+			ctx.ServerError("GetObsCreateSignedUrl", err)
+			return
+		}
+	} else if task.ComputeResource == models.GPUResource {
+		path := strings.TrimPrefix(path.Join(setting.TrainJobModelPath, task.JobName, setting.OutPutPath, versionName, parentDir, fileName), "/")
+		url, err = storage.GetObsCreateSignedUrlByBucketAndKey(setting.Bucket, path)
+		if err != nil {
+			log.Error("GetObsCreateSignedUrl failed: %v", err.Error(), ctx.Data["msgID"])
+			ctx.ServerError("GetObsCreateSignedUrl", err)
+			return
+		}
 	}
+
 	ctx.Resp.Header().Set("Cache-Control", "max-age=0")
 	http.Redirect(ctx.Resp, ctx.Req.Request, url, http.StatusMovedPermanently)
 }