From 2f9462965ebcca7695b8952eb801d8b9291b108b Mon Sep 17 00:00:00 2001
From: tzwang <wtz1003@sina.com>
Date: Tue, 12 Mar 2024 17:47:00 +0800
Subject: [PATCH] modified submit shuguangai implementation

Former-commit-id: 840111b6bb589bec04417f70e3d943f2fa999025
---
 api/internal/storeLink/shuguangai.go | 30 ++++++++++++++++++++++++++--
 api/internal/storeLink/storeLink.go  |  1 +
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/api/internal/storeLink/shuguangai.go b/api/internal/storeLink/shuguangai.go
index 242d07ec..36c9311b 100644
--- a/api/internal/storeLink/shuguangai.go
+++ b/api/internal/storeLink/shuguangai.go
@@ -352,15 +352,18 @@ func (s *ShuguangAi) generateResourceId(option *option.AiOption) error {
 
 	if option.ResourceType == CPU {
 		option.ResourceId = "WodTB2rJ8SobMgQ1nrtR245jxOrsovFi"
+		return nil
 	}
 
 	if option.ResourceType == CARD {
 		if option.Tops == 0 {
 			option.ResourceId = "WodTB2rJ8SobMgQ1nrtR245jxOrsovFi"
+			return nil
 		}
 
 		if option.Tops > DCU_TOPS {
 			option.ResourceId = "jeYBVPwyIALjVYNzHvysh2o5CsBpBLp2"
+			return nil
 		}
 
 		//Todo add more dcu specs
@@ -386,7 +389,12 @@ func (s *ShuguangAi) generateImageId(option *option.AiOption) error {
 		return errors.New("failed to get imageId")
 	}
 
-	if option.ResourceType == CPU {
+	for _, datum := range resp.Data {
+		ns := strings.Split(datum.Version, COLON)
+		if ns[0] == "jupyterlab-pytorch" {
+			option.ImageId = datum.ImageId
+			return nil
+		}
 
 	}
 
@@ -412,6 +420,7 @@ func (s *ShuguangAi) generateAlgorithmId(option *option.AiOption) error {
 		if ns[0] == option.DatasetsName {
 			algorithmId = option.TaskType + DASH + file.Name
 			option.AlgorithmId = algorithmId
+			option.AlgorithmName = ns[1]
 			return nil
 		}
 	}
@@ -430,6 +439,23 @@ func (s *ShuguangAi) generateEnv(option *option.AiOption) error {
 }
 
 func (s *ShuguangAi) generateParams(option *option.AiOption) error {
+	if option.ResourceType == "" {
+		return errors.New("ResourceType not set")
+	}
 
-	return nil
+	epoch := "epoch" + COMMA + "1"
+	option.Params = append(option.Params, epoch)
+
+	switch option.ResourceType {
+	case CPU:
+		card := "card" + COMMA + CPU
+		option.Params = append(option.Params, card)
+		return nil
+	case GPU:
+		card := "card" + COMMA + "cuda:0"
+		option.Params = append(option.Params, card)
+		return nil
+	}
+
+	return errors.New("failed to set params")
 }
diff --git a/api/internal/storeLink/storeLink.go b/api/internal/storeLink/storeLink.go
index 3693d58c..c77d3df2 100644
--- a/api/internal/storeLink/storeLink.go
+++ b/api/internal/storeLink/storeLink.go
@@ -40,6 +40,7 @@ type Linkage interface {
 }
 
 const (
+	COLON            = ":"
 	PY_PARAM_PREFIX  = "--"
 	SPACE            = " "
 	UNDERSCORE       = "_"