add models info file

largeModel/model_handler.go

@@ -8,19 +8,22 @@ import (
 	"example.com/m/operate"
 	"io"
 	"net/http"
+	"os"
 	"strings"
 	"time"
 )
 
 type ModelHandler struct {
-	dockerOp *operate.DockerOp
-	client   *http.Client
+	dockerOp       *operate.DockerOp
+	client         *http.Client
+	modelsFileName string
 }
 
 func NewModelHandler(dockerOp *operate.DockerOp) *ModelHandler {
 	return &ModelHandler{
-		dockerOp: dockerOp,
-		client:   &http.Client{},
+		dockerOp:       dockerOp,
+		client:         &http.Client{},
+		modelsFileName: "models.json",
 	}
 }
 
@@ -84,11 +87,38 @@ func (m *ModelHandler) MonitorModelInfo() {
 			}
 			m.dockerOp.ModelsInfo = modelInfosResp
 			m.dockerOp.ReportTaskIds = reportTaskIds
+			err = os.WriteFile(m.modelsFileName, bodyBytes, 0644)
+			if err != nil {
+				log.WithError(err).Error("Error writing models.json")
+			}
 			ticker = time.NewTicker(time.Minute * 10)
 		}
 	}
 }
 
+func (m *ModelHandler) ReadModels() ([]*models.ModelInfo, error) {
+	bodyBytes, err := os.ReadFile(m.modelsFileName)
+	if err != nil {
+		log.WithError(err).WithField("fileName", m.modelsFileName).Error("Error reading")
+		return nil, err
+	}
+	resp := &models.Resp{}
+	err = json.Unmarshal(bodyBytes, resp)
+	if err != nil {
+		log.WithField("fileName", m.modelsFileName).Error("Unmarshal model response failed:", err)
+		return nil, err
+	}
+	if resp.Code != http.StatusOK {
+		log.WithField("fileName", m.modelsFileName).Error("Response code :", resp.Code)
+		return nil, err
+	}
+	if resp.Data == nil || len(resp.Data) == 0 {
+		log.WithField("fileName", m.modelsFileName).Warn("Response data is empty")
+		return nil, err
+	}
+	return resp.Data, nil
+}
+
 func isResourceEnough(modelInfo *models.ModelInfo) bool {
 	return true
 }

nm/task_handler.go

@@ -132,7 +132,7 @@ func (t *TaskHandler) ComputeTaskHandler(taskMsg *nodeManagerV1.PushTaskMessage)
 		t.checkIsStopContainer(taskCmd)
 	}
 	log.Info("check is stop container finished")
-	isFound, imageId := t.foundTaskImage(taskExecResult, taskCmd, taskMsg)
+	isFound, imageId := t.foundTaskImage(taskCmd)
 	log.Info("found task image finished")
 	if !isFound || imageId == "" {
 		log.Error("The image is not found:", taskCmd.ImageName)
@@ -170,6 +170,7 @@ func (t *TaskHandler) ComputeTaskHandler(taskMsg *nodeManagerV1.PushTaskMessage)
 			isMatch := strings.HasPrefix(taskCmd.ImageName, models.ReplicateImageNameSuffix)
 			if isMatch {
 				if !t.checkContainerHealthy(internalIp, internalPort, taskMsg, taskExecResult) {
+					t.ExecTaskIdIsSuccess.Store(taskMsg.TaskId, true)
 					return
 				}
 			}
@@ -495,12 +496,10 @@ func (t *TaskHandler) CustomTaskHandler(taskMsg *nodeManagerV1.PushTaskMessage)
 	log.Info("received customTask--------------------------------")
 }
 
-func (t *TaskHandler) foundTaskImage(taskExecResult *models.TaskResult, taskCmd *models.TaskCmd, taskMsg *nodeManagerV1.PushTaskMessage) (isSuccess bool, imageId string) {
+func (t *TaskHandler) foundTaskImage(taskCmd *models.TaskCmd) (isSuccess bool, imageId string) {
 	images, err := t.DockerOp.PsImages()
 	if err != nil {
 		log.Error("Ps images failed:", err)
-		taskExecResult.TaskExecError = fmt.Sprintf("%s,%s", "Ps images failed:", err.Error())
-		t.ExecTaskIdIsSuccess.Store(taskMsg.TaskId, true)
 		isSuccess = false
 		imageId = ""
 		return
@@ -663,7 +662,6 @@ func (t *TaskHandler) checkContainerHealthy(internalIp string, internalPort uint
 	if err != nil {
 		log.Errorf("Request container healthy failed: %s", err.Error())
 		taskExecResult.TaskExecError = fmt.Sprintf("%s,%s", "Request container healthy failed", err.Error())
-		t.ExecTaskIdIsSuccess.Store(taskMsg.TaskId, true)
 		return false
 	}
 	if healthyCheckResp.StatusCode == http.StatusNotFound {
@@ -675,13 +673,11 @@ func (t *TaskHandler) checkContainerHealthy(internalIp string, internalPort uint
 	if err != nil {
 		log.Errorf("Json unmarshal container healthy body failed: %s", err.Error())
 		taskExecResult.TaskExecError = fmt.Sprintf("%s,%s", "Json unmarshal container healthy body failed", err.Error())
-		t.ExecTaskIdIsSuccess.Store(taskMsg.TaskId, true)
 		return false
 	}
 	if m.Status != models.READY {
 		log.Errorf("The container is not ready")
 		taskExecResult.TaskExecError = fmt.Sprintf("%s", "The container is not ready")
-		t.ExecTaskIdIsSuccess.Store(taskMsg.TaskId, true)
 		return false
 	}
 	return true