update code

db/workerinfo.go

@@ -121,9 +121,9 @@ func (d *dbWorker) FindWorkerByInstallModelAndSortByGpuRam(ctx context.Context,
 	// sort by gpu ram
 	findOptions := options.Find()
 	findOptions.SetLimit(int64(limit))
-	findOptions.SetSort(bson.D{{"hardware.gpu.ram", 1}})
+	//findOptions.SetSort(bson.D{{"hardware.GPU.mem_free", 1}})
 
-	selector := bson.M{"model_infos.installed_models.model_id": modelId, "hardware.gpu.performance": bson.M{"$gte": performance}, "hardware.gpu.ram": bson.M{"$gte": ram}}
+	selector := bson.M{"model_infos.installed_models.model_id": modelId, "hardware.GPU.mem_free": bson.M{"$gte": ram}}
 	cursor, err := d.col.Find(ctx, selector, findOptions)
 	if err != nil {
 		return nil, err

db/workerinfo_test.go

@@ -271,6 +271,7 @@ func TestDbWorker_InsertWorker(t *testing.T) {
 
 	wg := sync.WaitGroup{}
 	for i := 0; i < 5; i++ {
+		wg.Add(1)
 		go func() {
 			defer wg.Done()
 			for {
@@ -455,7 +456,7 @@ func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam(b *testing.B) {
 	defer db.client.Disconnect(context.Background())
 	b.ResetTimer()
 	for i := 0; i < b.N; i++ {
-		installedModelId := getRandId(maxModelId)
+		installedModelId := getRandId(100)
 		performance := generateAGpuPerformance()
 		ram := generateAGpuRam()
 		if w, err := db.FindWorkerByInstallModelAndSortByGpuRam(context.Background(), installedModelId, performance, ram, 10); err != nil {
@@ -475,7 +476,7 @@ func BenchmarkDbWorker_FindWorkerByInstallModelAndSortByGpuRam_Parallel(b *testi
 	defer db.client.Disconnect(context.Background())
 	b.RunParallel(func(pb *testing.PB) {
 		for pb.Next() {
-			installedModelId := getRandId(maxModelId)
+			installedModelId := getRandId(100)
 			performance := generateAGpuPerformance()
 			ram := generateAGpuRam()
 			if w, err := db.FindWorkerByInstallModelAndSortByGpuRam(context.Background(), installedModelId, performance, ram, 10); err != nil {