update worker info and test

db/workerinfo.go

@@ -95,28 +95,12 @@ func (d *dbWorker) UpdateNodeInfo(ctx context.Context, id string, nodeInfo *type
 	return err
 }
 
-func (d *dbWorker) FindWorkerByInstalledModelId(ctx context.Context, modelId string) ([]*DbWorkerInfo, error) {
-	// find all worker that at least one installed model's mode_id is equal modelId
-	selector := bson.M{"model_infos.installed_models.model_id": modelId}
-
-	cursor, err := d.col.Find(ctx, selector)
-	if err != nil {
-		return nil, err
-	}
-	defer cursor.Close(ctx)
-
-	var workers []*DbWorkerInfo
-	if err = cursor.All(ctx, &workers); err != nil {
-		return nil, err
-	}
-	return workers, nil
-}
-
-func (d *dbWorker) FindWorkerByRunningModelIdWithLimit(ctx context.Context, modelId string, limit int64) ([]*DbWorkerInfo, error) {
+func (d *dbWorker) FindWorkerByRunningModelAndSortByWaitTime(ctx context.Context, modelId string, limit int) ([]*DbWorkerInfo, error) {
 	// find all worker that at least one running model's mode_id is equal modelId
+	// sort by wait time
 	findOptions := options.Find()
-	findOptions.SetLimit(limit)
-	findOptions.SetSort(bson.D{{"hardware.gpu.usage", 1}})
+	findOptions.SetLimit(int64(limit))
+	findOptions.SetSort(bson.D{{"model_infos.running_models.wait_time", 1}})
 
 	selector := bson.M{"model_infos.running_models.model_id": modelId}
 	cursor, err := d.col.Find(ctx, selector, findOptions)
@@ -132,27 +116,15 @@ func (d *dbWorker) FindWorkerByRunningModelIdWithLimit(ctx context.Context, mode
 	return workers, nil
 }
 
-func (d *dbWorker) FindWorkerByRunningModelId(ctx context.Context, modelId string) ([]*DbWorkerInfo, error) {
-	// find all worker that at least one running model's mode_id is equal modelId
-	selector := bson.M{"model_infos.running_models.model_id": modelId}
-	cursor, err := d.col.Find(ctx, selector)
-	if err != nil {
-		return nil, err
-	}
-	defer cursor.Close(ctx)
-
-	var workers []*DbWorkerInfo
-	if err = cursor.All(ctx, &workers); err != nil {
-		return nil, err
-	}
-	return workers, nil
-}
-
-func (d *dbWorker) FindWorkerByGpuModel(ctx context.Context, model string) ([]*DbWorkerInfo, error) {
-	// find all worker that at least one gpu model is equal model.
-	selector := bson.M{"hardware.gpu.model": model}
+func (d *dbWorker) FindWorkerByInstallModelAndSortByGpuRam(ctx context.Context, modelId string, performance int, ram int, limit int) ([]*DbWorkerInfo, error) {
+	// find all worker that at least one installed model's mode_id is equal modelId
+	// sort by gpu ram
+	findOptions := options.Find()
+	findOptions.SetLimit(int64(limit))
+	findOptions.SetSort(bson.D{{"hardware.gpu.ram", 1}})
 
-	cursor, err := d.col.Find(ctx, selector)
+	selector := bson.M{"model_infos.installed_models.model_id": modelId, "hardware.gpu.performance": bson.M{"$gte": performance}, "hardware.gpu.ram": bson.M{"$gte": ram}}
+	cursor, err := d.col.Find(ctx, selector, findOptions)
 	if err != nil {
 		return nil, err
 	}
@@ -163,21 +135,5 @@ func (d *dbWorker) FindWorkerByGpuModel(ctx context.Context, model string) ([]*D
 		return nil, err
 	}
 	return workers, nil
-}
-
-func (d *dbWorker) FindWorkerByGpuRam(ctx context.Context, ram uint) ([]*DbWorkerInfo, error) {
-	// find all worker that at least one gpu ram is greater or equal than ram.
-	selector := bson.M{"hardware.gpu.ram": bson.M{"$gte": ram}}
 
-	cursor, err := d.col.Find(ctx, selector)
-	if err != nil {
-		return nil, err
-	}
-	defer cursor.Close(ctx)
-
-	var workers []*DbWorkerInfo
-	if err = cursor.All(ctx, &workers); err != nil {
-		return nil, err
-	}
-	return workers, nil
 }

types/worker.go

@@ -6,65 +6,71 @@ type NodeInfo struct {
 	DeviceIp       string `bson:"device_ip" json:"device_ip"`
 }
 
-type InstalledModelInfo struct {
-	ModelId  string `bson:"model_id" json:"model_id"`
-	DiskSize uint64 `bson:"disk_size" json:"disk_size"`
+type InstalledModel struct {
+	ModelID       string `bson:"model_id" json:"model_id"`
+	DiskSize      int    `bson:"disk_size" json:"disk_size"`
+	InstalledTime int64  `bson:"installed_time" json:"installed_time"`
+	LastRunTime   int64  `bson:"last_run_time" json:"last_run_time"`
 }
 
-type RunningModelInfo struct {
-	ModelId     string `bson:"model_id" json:"model_id"`
-	StartedTime uint64 `bson:"started_time" json:"started_time"`
-	LatestTime  uint64 `bson:"latest_time" json:"latest_time"`
-	RunCount    uint64 `bson:"run_count" json:"run_count"`
-	GpuRamUsed  uint64 `bson:"gpu_ram" json:"gpu_ram"`
+type RunningModel struct {
+	ModelID       string `bson:"model_id" json:"model_id"`
+	GpuSeq        int    `bson:"gpu_seq" json:"gpu_seq"`
+	GpuRAM        int    `bson:"gpu_ram" json:"gpu_ram"`
+	StartedTime   int64  `bson:"started_time" json:"started_time"`
+	LastWorkTime  int64  `bson:"last_work_time" json:"last_work_time"`
+	TotalRunCount int    `bson:"total_run_count" json:"total_run_count"`
+	WaitTime      int    `bson:"wait_time" json:"wait_time"`
 }
 
 type ModelInfo struct {
-	InstalledModels []*InstalledModelInfo `bson:"installed_models" json:"installed_models"`
-	RunningModels   []*RunningModelInfo   `bson:"running_models" json:"running_models"`
-}
-
-type DeviceInfo struct {
-	DeviceType  string `bson:"device_type" json:"device_type"`
-	DeviceModel string `bson:"device_model" json:"device_model"`
-	DeviceParam string `bson:"device_param" json:"device_param"`
-	DevicePower uint64 `bson:"device_power" json:"device_power"`
+	InstalledModels []*InstalledModel `bson:"installed_models" json:"installed_models"`
+	RunningModels   []*RunningModel   `bson:"running_models" json:"running_models"`
 }
 
-type DeviceUsage struct {
-	DeviceType  string `bson:"device_type" json:"device_type"`
-	DeviceUsage uint64 `bson:"device_usage" json:"device_usage"`
+type HardwareInfo struct {
+	CPU  *CpuInfo   `bson:"CPU" json:"CPU"`
+	GPU  []*GpuInfo `bson:"GPU" json:"GPU"`
+	RAM  *RamInfo   `bson:"RAM" json:"RAM"`
+	DISK *DiskInfo  `bson:"DISK" json:"DISK"`
+	NET  *NetInfo   `bson:"NET" json:"NET"`
 }
 
 type CpuInfo struct {
-	Model string `bson:"model" json:"model"`
-	Core  uint   `bson:"core" json:"core"`
-	Usage uint   `bson:"usage" json:"usage"`
+	Model   string `bson:"model" json:"model "`
+	Number  int    `bson:"number" json:"number"`
+	Cores   int    `bson:"cores" json:"cores"`
+	Threads int    `bson:"threads" json:"threads"`
+	Usage   int    `bson:"usage" json:"usage"`
 }
 
 type GpuInfo struct {
-	Model  string `bson:"model" json:"model"`
-	Ram    uint   `bson:"ram" json:"ram"`
-	Usage  uint   `bson:"usage" json:"usage"`
-	Occupy uint   `bson:"occupy" json:"occupy"`
-	Temp   uint   `bson:"temp" json:"temp"`
+	Seq         int    `bson:"seq" json:"seq"`
+	UUID        string `bson:"uuid" json:"uuid"`
+	Model       string `bson:"model" json:"model "`
+	Performance int    `bson:"performance" json:"performance"`
+	PowerRating int    `bson:"power_rating" json:"power_rating"`
+	MemTotal    int    `bson:"mem_total" json:"mem_total"`
+	MemFree     int    `bson:"mem_free" json:"mem_free"`
+	Usage       int    `bson:"usage" json:"usage"`
+	Temp        int    `bson:"temp" json:"temp "`
+	PowerRt     int    `bson:"power_rt" json:"power_rt"`
 }
 
 type RamInfo struct {
-	Size  uint `bson:"size" json:"size"`
-	Usage uint `bson:"usage" json:"usage"`
+	Total int `bson:"total" json:"total"`
+	Free  int `bson:"free" json:"free"`
 }
 
 type DiskInfo struct {
-	Size  uint `bson:"size" json:"size"`
-	Usage uint `bson:"usage" json:"usage"`
+	Total int `bson:"total" json:"total"`
+	Free  int `bson:"free" json:"free"`
 }
 
-type HardwareInfo struct {
-	Cpu  []*CpuInfo `bson:"cpu" json:"cpu"`
-	Gpu  []*GpuInfo `bson:"gpu" json:"gpu"`
-	Ram  *RamInfo   `bson:"ram" json:"ram"`
-	Disk *DiskInfo  `bson:"disk" json:"disk"`
+type NetInfo struct {
+	IP        string `bson:"ip" json:"ip"`
+	Mac       string `bson:"mac" json:"mac"`
+	Bandwidth int    `bson:"bandwidth" json:"bandwidth"`
 }
 
 type WorkerInfo struct {
@@ -73,3 +79,8 @@ type WorkerInfo struct {
 	ModelInofs *ModelInfo    `bson:"model_infos" json:"model_infos"`
 	Hardware   *HardwareInfo `bson:"hardware" json:"hardware"`
 }
+
+type WorkerModelInfo struct {
+	WorkerId string `bson:"worker_id" json:"worker_id"`
+	ModelId  string `bson:"model_id" json:"model_id"`
+}