自定义镜像支持。

Signed-off-by: zouap <zouap@pcl.ac.cn>

modules/modelarts/modelarts.go

@@ -138,6 +138,7 @@ type VersionInfo struct {
 	Version []struct {
 		ID    int    `json:"id"`
 		Value string `json:"value"`
+		Url   string `json:"url"`
 	} `json:"version"`
 }
 
@@ -152,7 +153,6 @@ type Engine struct {
 	Info []struct {
 		ID    int    `json:"id"`
 		Value string `json:"value"`
-		Url   string `json:"url"`
 	} `json:"engine"`
 }
 

routers/repo/modelarts.go

@@ -1216,14 +1216,15 @@ func getUserCommand(engineId int, req *modelarts.GenerateTrainJobReq) (string, s
 	userCommand := ""
 	if engineId < 0 {
 		userCommand = "/bin/bash /home/work/run_train.sh 's3://" + req.CodeObsPath + "' 'code/" + req.BootFile + "' '/tmp/log/train.log' --'data_url'='s3://" + req.DataUrl + "' --'train_url'='s3://" + req.TrainUrl + "'"
-		var engines modelarts.Engine
-		if err := json.Unmarshal([]byte(setting.Engines), &engines); err != nil {
-			log.Info("json.Unmarshal failed:", err)
-		}
-		for _, engine := range engines.Info {
-			if engine.ID == engineId {
-				userImageUrl = engine.Url
-				break
+		var versionInfos modelarts.VersionInfo
+		if err := json.Unmarshal([]byte(setting.EngineVersions), &versionInfos); err != nil {
+			log.Info("json parse err." + err.Error())
+		} else {
+			for _, engine := range versionInfos.Version {
+				if engine.ID == engineId {
+					userImageUrl = engine.Url
+					break
+				}
 			}
 		}
 		for _, param := range req.Parameters {