update test

operator/workerInstalledInfo.go

@@ -46,7 +46,7 @@ func (d *WorkerInstalledOperator) FindWorkerByModelId(ctx context.Context, model
 	// sort by wait time
 	findOptions := options.Find()
 	findOptions.SetLimit(int64(limit))
-	findOptions.SetSort(bson.D{{"exec_time", 1}})
+	findOptions.SetSort(bson.D{{"gpu_free", 1}})
 
 	selector := bson.M{"model_id": modelId}
 	cursor, err := d.col.Find(ctx, selector, findOptions)
@@ -61,3 +61,24 @@ func (d *WorkerInstalledOperator) FindWorkerByModelId(ctx context.Context, model
 	}
 	return workers, nil
 }
+
+func (d *WorkerInstalledOperator) FindWorkerByModelIdAndGpuMem(ctx context.Context, modelId int, mem int64, limit int) ([]*WorkerInstalledInfo, error) {
+	// find all worker that at least one installed model's mode_id is equal modelId
+	// sort by gpu_free
+	findOptions := options.Find()
+	findOptions.SetLimit(int64(limit))
+	findOptions.SetSort(bson.D{{"gpu_free", 1}})
+
+	selector := bson.M{"model_id": modelId, "gpu_free": bson.M{"$gte": mem}}
+	cursor, err := d.col.Find(ctx, selector, findOptions)
+	if err != nil {
+		return nil, err
+	}
+	defer cursor.Close(ctx)
+
+	var workers []*WorkerInstalledInfo
+	if err = cursor.All(ctx, &workers); err != nil {
+		return nil, err
+	}
+	return workers, nil
+}

operator/workerInstalledInfo_test.go

@@ -105,6 +105,83 @@ func BenchmarkWorkerInstalledOperator_UpdateGpuFree_Parallel(b *testing.B) {
 	})
 }
 
+func BenchmarkWorkerInstalledOperator_FindWorkerByModelId(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		id := rand.Intn(maxModelId)
+		_, err := db.FindWorkerByModelId(context.Background(), id, 10)
+		if err != nil {
+			panic(fmt.Sprintf("find worker failed with err:%s", err))
+		}
+	}
+
+}
+
+func BenchmarkWorkerInstalledOperator_FindWorkerByModelId_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			id := rand.Intn(maxModelId)
+			_, err := db.FindWorkerByModelId(context.Background(), id, 10)
+			if err != nil {
+				panic(fmt.Sprintf("find worker failed with err:%s", err))
+			}
+		}
+	})
+
+}
+
+func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		id := rand.Intn(maxModelId)
+		mem := rand.Intn(100) + 12000
+		_, err := db.FindWorkerByModelIdAndGpuMem(context.Background(), id, int64(mem), 10)
+		if err != nil {
+			panic(fmt.Sprintf("find worker failed with err:%s", err))
+		}
+	}
+}
+
+func BenchmarkWorkerInstalledOperator_FindWorkerByModelIdAndGpuMem_Parallel(b *testing.B) {
+	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
+	if err != nil {
+		log.Fatal(err)
+	}
+	db := NewDBWorkerInstalled(client, database)
+	defer db.client.Disconnect(context.Background())
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			id := rand.Intn(maxModelId)
+			mem := rand.Intn(100) + 12000
+			_, err := db.FindWorkerByModelIdAndGpuMem(context.Background(), id, int64(mem), 10)
+			if err != nil {
+				panic(fmt.Sprintf("find worker failed with err:%s", err))
+			}
+		}
+	})
+
+}
+
 func TestWorkerInstalledOperator_UpdateGpuFree(t *testing.T) {
 	client, err := ConnectMongoDB("mongodb://localhost:27017", "admin", "admin")
 	if err != nil {