change hardware and add test

db/workerinfo.go

@@ -93,3 +93,70 @@ func (d *dbWorker) UpdateNodeInfo(ctx context.Context, id string, nodeInfo *type
 	_, err := d.col.UpdateOne(ctx, bson.M{"_id": id}, update)
 	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) 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}
+
+	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) 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
+}

db/workerinfo_test.go

@@ -83,25 +83,66 @@ func generateANodeInfo() *types.NodeInfo {
 	}
 }
 
+func generateACpu() *types.CpuInfo {
+	return &types.CpuInfo{
+		Model: "Intel i9",
+		Core:  8,
+		Usage: uint(rand.Intn(50) + 30),
+	}
+}
+
+func generateADisk() *types.DiskInfo {
+	return &types.DiskInfo{
+		Size:  1024,
+		Usage: uint(rand.Intn(100) + 100),
+	}
+}
+
+func generateARam() *types.RamInfo {
+	return &types.RamInfo{
+		Size:  8 * (1 << 30),
+		Usage: uint(rand.Intn(30) + 40),
+	}
+}
+
+func generateAGpuRam() uint {
+	return uint(rand.Intn(3)*8 + 8) // 8, 16, 24
+}
+
+func generateAGpuModel() string {
+	m := rand.Intn(4)*10 + 3060 // 3060 ~ 3090
+	return fmt.Sprintf("Nvidia %d", m)
+}
+
+func generateAIdleGpu() *types.GpuInfo {
+	return &types.GpuInfo{
+		Model: generateAGpuModel(),
+		Ram:   generateAGpuRam(),
+	}
+}
+
+func generateAUsageGpu() *types.GpuInfo {
+	ram := generateAGpuRam()
+	return &types.GpuInfo{
+		Model:  generateAGpuModel(),
+		Ram:    ram,
+		Usage:  uint(rand.Intn(10) + 30),
+		Occupy: ram * 10 / 7,
+		Temp:   60,
+	}
+}
+
 func generateAHardware() *types.HardwareInfo {
 	return &types.HardwareInfo{
-		Devices: []*types.DeviceInfo{
-			&types.DeviceInfo{
-				DeviceType:  "cpu",
-				DeviceModel: "xxxxxxxxx",
-				DeviceParam: "kkkkkkkk",
-				DevicePower: 0,
-			},
-			&types.DeviceInfo{
-				DeviceType:  "gpu",
-				DeviceModel: "Nvidia 3090",
-				DeviceParam: "kkkkkkkk",
-				DevicePower: 100,
-			},
+		Cpu: []*types.CpuInfo{
+			generateACpu(),
 		},
-		Usages: []*types.DeviceUsage{
-			&types.DeviceUsage{},
+		Gpu: []*types.GpuInfo{
+			generateAIdleGpu(),
+			generateAUsageGpu(),
 		},
+		Ram:  generateARam(),
+		Disk: generateADisk(),
 	}
 }
 
@@ -284,3 +325,114 @@ func BenchmarkDbWorker_UpdateNodeInfo_Parallel(b *testing.B) {
 		}
 	})
 }
+
+func BenchmarkDbWorker_FindWorkerByRunningModelId(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorker(client, database, collection)
+	defer db.client.Disconnect(context.Background())
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		runningModelId := getRandId(100)
+		if w, err := db.FindWorkerByRunningModelId(context.Background(), runningModelId); err != nil {
+			panic(fmt.Sprintf("find worker failed with err:%s", err))
+		} else if len(w) == 0 {
+			b.Logf("FindWorkerByRunningModelId find %d with id %d\n", len(w), runningModelId)
+		}
+	}
+}
+
+func BenchmarkDbWorker_FindWorkerByRunningModelId_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorker(client, database, collection)
+	defer db.client.Disconnect(context.Background())
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			runningModelId := getRandId(100)
+			if w, err := db.FindWorkerByRunningModelId(context.Background(), runningModelId); err != nil {
+				panic(fmt.Sprintf("find worker failed with err:%s", err))
+			} else if len(w) == 0 {
+				b.Logf("FindWorkerByRunningModelId find %d with id %d\n", len(w), runningModelId)
+			}
+		}
+	})
+}
+
+func BenchmarkDbWorker_FindWorkerByInstalledModelId(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorker(client, database, collection)
+	defer db.client.Disconnect(context.Background())
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		installedModelId := getRandId(100)
+		if w, err := db.FindWorkerByInstalledModelId(context.Background(), installedModelId); err != nil {
+			panic(fmt.Sprintf("find worker failed with err:%s", err))
+		} else if len(w) == 0 {
+			b.Logf("FindWorkerByInstalledModelId find %d with id %d\n", len(w), installedModelId)
+		}
+	}
+}
+
+func BenchmarkDbWorker_FindWorkerByInstalledModelId_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorker(client, database, collection)
+	defer db.client.Disconnect(context.Background())
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			installedModelId := getRandId(100)
+			if w, err := db.FindWorkerByInstalledModelId(context.Background(), installedModelId); err != nil {
+				panic(fmt.Sprintf("find worker failed with err:%s", err))
+			} else if len(w) == 0 {
+				b.Logf("FindWorkerByInstalledModelId find %d with id %d\n", len(w), installedModelId)
+			}
+		}
+	})
+}
+
+func BenchmarkDbWorker_FindWorkerByGpuRam(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorker(client, database, collection)
+	defer db.client.Disconnect(context.Background())
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		ram := generateAGpuRam()
+		if w, err := db.FindWorkerByGpuRam(context.Background(), ram); err != nil {
+			panic(fmt.Sprintf("find worker failed with err:%s", err))
+		} else if len(w) == 0 {
+			b.Logf("FindWorkerByGpuRam find %d with ram %d\n", len(w), ram)
+		}
+	}
+}
+
+func BenchmarkDbWorker_FindWorkerByGpuRam_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorker(client, database, collection)
+	defer db.client.Disconnect(context.Background())
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			ram := generateAGpuRam()
+			if w, err := db.FindWorkerByGpuRam(context.Background(), ram); err != nil {
+				panic(fmt.Sprintf("find worker failed with err:%s", err))
+			} else if len(w) == 0 {
+				b.Logf("FindWorkerByGpuRam find %d with ram %d\n", len(w), ram)
+			}
+		}
+	})
+}

types/worker.go

@@ -36,9 +36,35 @@ type DeviceUsage struct {
 	DeviceUsage uint64 `bson:"device_usage" json:"device_usage"`
 }
 
+type CpuInfo struct {
+	Model string `bson:"model" json:"model"`
+	Core  uint   `bson:"core" json:"core"`
+	Usage uint   `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"`
+}
+
+type RamInfo struct {
+	Size  uint `bson:"size" json:"size"`
+	Usage uint `bson:"usage" json:"usage"`
+}
+
+type DiskInfo struct {
+	Size  uint `bson:"size" json:"size"`
+	Usage uint `bson:"usage" json:"usage"`
+}
+
 type HardwareInfo struct {
-	Devices []*DeviceInfo  `bson:"devices" json:"devices"`
-	Usages  []*DeviceUsage `bson:"usages" json:"usages"`
+	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 WorkerInfo struct {