add new case

db/workerinfo.go

@@ -6,6 +6,7 @@ import (
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/primitive"
 	"go.mongodb.org/mongo-driver/mongo"
+	"go.mongodb.org/mongo-driver/mongo/options"
 	"mogo/types"
 )
 
@@ -111,6 +112,26 @@ func (d *dbWorker) FindWorkerByInstalledModelId(ctx context.Context, modelId str
 	return workers, nil
 }
 
+func (d *dbWorker) FindWorkerByRunningModelIdWithLimit(ctx context.Context, modelId string, limit int64) ([]*DbWorkerInfo, error) {
+	// find all worker that at least one running model's mode_id is equal modelId
+	findOptions := options.Find()
+	findOptions.SetLimit(limit)
+	findOptions.SetSort(bson.D{{"hardware.gpu.usage", 1}})
+
+	selector := bson.M{"model_infos.running_models.model_id": modelId}
+	cursor, err := d.col.Find(ctx, selector, findOptions)
+	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}

db/workerinfo_test.go

@@ -363,6 +363,43 @@ func BenchmarkDbWorker_FindWorkerByRunningModelId_Parallel(b *testing.B) {
 	})
 }
 
+func BenchmarkDbWorker_FindWorkerByRunningModelIdWithLimit(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.FindWorkerByRunningModelIdWithLimit(context.Background(), runningModelId, 10); err != nil {
+			panic(fmt.Sprintf("find worker failed with err:%s", err))
+		} else if len(w) == 0 {
+			b.Logf("FindWorkerByRunningModelId find %d with id %s\n", len(w), runningModelId)
+		}
+	}
+}
+
+func BenchmarkDbWorker_FindWorkerByRunningModelIdWithLimit_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.FindWorkerByRunningModelIdWithLimit(context.Background(), runningModelId, 10); err != nil {
+				panic(fmt.Sprintf("find worker failed with err:%s", err))
+			} else if len(w) == 0 {
+				b.Logf("FindWorkerByRunningModelId find %d with id %s\n", len(w), runningModelId)
+			}
+		}
+	})
+}
+
 func BenchmarkDbWorker_FindWorkerByInstalledModelId(b *testing.B) {
 	client, err := ConnectMongoDB("mongodb://localhost:27017")
 	if err != nil {