update db and msg resp

db/db.go

@@ -22,6 +22,15 @@ func init() {
 		log.Error("Leveldb open file failed: ", err)
 	}
 	modelKeys = make(map[string]bool, 0)
+	iter := dbInstance.NewIterator(nil, nil)
+	defer iter.Release()
+	for iter.Next() {
+		key := iter.Key()
+		modelKeys[string(key)] = true
+	}
+	if err := iter.Error(); err != nil {
+		log.Fatal(err)
+	}
 }
 
 func Put(key string, value any) error {

largeModel/model_handler.go

@@ -66,10 +66,11 @@ func (m *ModelHandler) MonitorModelInfo() {
 			//	continue
 			//}
 			modelInfosResp := resp.Data
-			for _, modelInfo := range modelInfosResp {
+			for i, modelInfo := range modelInfosResp {
 				if modelInfo.ImageName == "" {
 					continue
 				}
+
 				modelInfo.ImageName = fmt.Sprintf("%s-%s", modelInfo.ImageName, conf.GetConfig().OpSys)
 				split := strings.Split(modelInfo.ImageName, ":")
 				if len(split) != 2 {
@@ -79,6 +80,19 @@ func (m *ModelHandler) MonitorModelInfo() {
 				m.dockerOp.SignApi[modelInfo.ImageName] = modelInfo.SignUrl
 				model, _ := db.GetModel(modelInfo.ImageName)
 				if model != nil {
+					switch i {
+					case 0:
+						model.IsInstalled = true
+						model.IsRunning = true
+						model.GpuSeq = 0
+					case 1:
+						model.IsInstalled = true
+						model.IsRunning = true
+						model.GpuSeq = 1
+					default:
+						model.IsInstalled = true
+						model.GpuSeq = 999
+					}
 					model.UpdateFiled(modelInfo)
 					err := db.PutModel(modelInfo.ImageName, model)
 					if err != nil {
@@ -87,6 +101,19 @@ func (m *ModelHandler) MonitorModelInfo() {
 					}
 					log.WithField("name", modelInfo.ImageName).Info("The image updated")
 				} else {
+					switch i {
+					case 0:
+						modelInfo.IsInstalled = true
+						modelInfo.IsRunning = true
+						modelInfo.GpuSeq = 0
+					case 1:
+						modelInfo.IsInstalled = true
+						modelInfo.IsRunning = true
+						modelInfo.GpuSeq = 1
+					default:
+						modelInfo.IsInstalled = true
+						modelInfo.GpuSeq = 999
+					}
 					err := db.PutModel(modelInfo.ImageName, modelInfo)
 					if err != nil {
 						log.WithError(err).Error("Put db error")
@@ -211,32 +238,32 @@ func (m *ModelHandler) ScanModelsResp() (*nodemanagerV2.ModelsInfo, error) {
 	//	log.Error("Get container failed")
 	//	return nil, fmt.Errorf("get containe failed")
 	//}
-	//allModels, err := db.GetAllModels()
-	//if err != nil {
-	//	log.WithError(err).Error("Get all models failed")
-	//	return nil, fmt.Errorf("get all models failed")
-	//}
-	//for _, model := range allModels {
-	//	isExist := images[model.ImageName]
-	//	if !isExist {
-	//		continue
-	//	}
-	//	diskSize, err := strconv.ParseInt(model.HardwareRequire.DiskSize, 10, 64)
-	//	if err != nil {
-	//		continue
-	//	}
-	//	installedModels = append(installedModels, &nodemanagerV2.InstalledModel{ModelId: strconv.FormatUint(model.TaskId, 10), DiskSize: diskSize, InstalledTime: model.SetupTime, LastRunTime: model.LastRunTime})
-	//
-	//	containerIsExist := false
-	//	for _, container := range containerList {
-	//		if model.ImageName == container.Image {
-	//			containerIsExist = true
-	//		}
-	//	}
-	//	if containerIsExist {
-	//		runningModels = append(runningModels, &nodemanagerV2.RunningModel{ModelId: strconv.FormatUint(model.TaskId, 10), GpuSeq: model.GpuSeq, GpuRam: model.RunningMem, StartedTime: model.LastRunTime, LastWorkTime: model.LastWorkTime, TotalRunCount: model.TotalRunCount, ExecTime: model.EstimatExeTime})
-	//	}
-	//}
+	allModels, err := db.GetAllModels()
+	if err != nil {
+		log.WithError(err).Error("Get all models failed")
+		return nil, fmt.Errorf("get all models failed")
+	}
+	for _, model := range allModels {
+		//isExist := images[model.ImageName]
+		//if !isExist {
+		//	continue
+		//}
+		if model.IsInstalled {
+			diskSize, err := strconv.ParseInt(model.HardwareRequire.DiskSize, 10, 64)
+			if err != nil {
+				continue
+			}
+			installedModels = append(installedModels, &nodemanagerV2.InstalledModel{ModelId: strconv.FormatUint(model.TaskId, 10), DiskSize: diskSize, InstalledTime: model.SetupTime, LastRunTime: model.LastRunTime})
+		}
+		//for _, container := range containerList {
+		//	if model.ImageName == container.Image {
+		//		containerIsExist = true
+		//	}
+		//}
+		if model.IsRunning {
+			runningModels = append(runningModels, &nodemanagerV2.RunningModel{ModelId: strconv.FormatUint(model.TaskId, 10), GpuSeq: model.GpuSeq, GpuRam: model.RunningMem, StartedTime: model.LastRunTime, LastWorkTime: model.LastWorkTime, TotalRunCount: model.TotalRunCount, ExecTime: model.EstimatExeTime})
+		}
+	}
 	res := &nodemanagerV2.ModelsInfo{
 		InstalledModels: installedModels,
 		RunningModels:   runningModels,

nm/start.go

@@ -42,7 +42,10 @@ func StartMonitor() {
 	go monitorNm.monitorNodeManagerSeed()
 	log.WithField("func", "monitorNodeManagerSeed").Info("--------------------Start monitorNm--------------------")
 
-	for !monitorNm.IsInit && !modelHandler.IsInit {
+	for {
+		if monitorNm.IsInit && modelHandler.IsInit {
+			break
+		}
 		time.Sleep(time.Second * 3)
 	}