profile

operator/workerinfo.go

@@ -205,7 +205,11 @@ func (d *WorkerInfoOperator) UpdateGPUUsage(ctx context.Context, workerid string
 	return err
 }
 
-func (d *WorkerInfoOperator) FindWorkerByRunningModelAndSortByWaitTime(ctx context.Context, modelId string, limit int) ([]*WorkerInfo, error) {
+type WorkerFindResult struct {
+	WorkerId string `bson:"worker_id" json:"worker_id"`
+}
+
+func (d *WorkerInfoOperator) FindWorkerByRunningModelAndSortByWaitTime(ctx context.Context, modelId string, limit int) ([]*WorkerFindResult, error) {
 	// find all worker that at least one running model's mode_id is equal modelId
 	// sort by wait time
 	findOptions := options.Find()
@@ -219,14 +223,14 @@ func (d *WorkerInfoOperator) FindWorkerByRunningModelAndSortByWaitTime(ctx conte
 	}
 	defer cursor.Close(ctx)
 
-	var workers []*WorkerInfo
-	if err = cursor.All(ctx, &workers); err != nil {
+	var find []*WorkerFindResult
+	if err = cursor.All(ctx, &find); err != nil {
 		return nil, err
 	}
-	return workers, nil
+	return find, nil
 }
 
-func (d *WorkerInfoOperator) FindWorkerByInstallModelAndSortByGpuRam(ctx context.Context, modelId string, performance int, ram int64, limit int) ([]*WorkerInfo, error) {
+func (d *WorkerInfoOperator) FindWorkerByInstallModelAndSortByGpuRam(ctx context.Context, modelId string, performance int, ram int64, limit int) ([]*WorkerFindResult, error) {
 	// find all worker that at least one installed model's mode_id is equal modelId
 	// sort by gpu ram
 	findOptions := options.Find()
@@ -240,11 +244,11 @@ func (d *WorkerInfoOperator) FindWorkerByInstallModelAndSortByGpuRam(ctx context
 	}
 	defer cursor.Close(ctx)
 
-	var workers []*WorkerInfo
-	if err = cursor.All(ctx, &workers); err != nil {
+	var find []*WorkerFindResult
+	if err = cursor.All(ctx, &find); err != nil {
 		return nil, err
 	}
-	return workers, nil
+	return find, nil
 }
 
 func (d *WorkerInfoOperator) FindWorkerByWorkerId(ctx context.Context, workerId string) (*WorkerInfo, error) {