milvus/internal/util/function/embedding/function_executor.go

/*
 * # Licensed to the LF AI & Data foundation under one
 * # or more contributor license agreements. See the NOTICE file
 * # distributed with this work for additional information
 * # regarding copyright ownership. The ASF licenses this file
 * # to you under the Apache License, Version 2.0 (the
 * # "License"); you may not use this file except in compliance
 * # with the License. You may obtain a copy of the License at
 * #
 * #     http://www.apache.org/licenses/LICENSE-2.0
 * #
 * # Unless required by applicable law or agreed to in writing, software
 * # distributed under the License is distributed on an "AS IS" BASIS,
 * # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * # See the License for the specific language governing permissions and
 * # limitations under the License.
 */

package embedding

import (
	"context"
	"strconv"
	"sync"

	"google.golang.org/protobuf/proto"

	"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
	"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
	"github.com/milvus-io/milvus/internal/storage"
	"github.com/milvus-io/milvus/internal/util/function"
	"github.com/milvus-io/milvus/internal/util/function/models"
	"github.com/milvus-io/milvus/pkg/v3/metrics"
	"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
	"github.com/milvus-io/milvus/pkg/v3/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v3/util/merr"
	"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
	"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
)

type Runner interface {
	GetSchema() *schemapb.FunctionSchema
	GetOutputFields() []*schemapb.FieldSchema
	GetCollectionName() string
	GetFunctionTypeName() string
	GetFunctionName() string
	GetFunctionProvider() string
	Check(ctx context.Context) error

	MaxBatch() int
	ProcessInsert(ctx context.Context, inputs []*schemapb.FieldData) ([]*schemapb.FieldData, error)
	ProcessSearch(ctx context.Context, placeholderGroup *commonpb.PlaceholderGroup) (*commonpb.PlaceholderGroup, error)
	ProcessBulkInsert(ctx context.Context, inputs []storage.FieldData) (map[storage.FieldID]storage.FieldData, error)
}

type FunctionExecutor struct {
	runners map[int64]Runner
}

func createFunction(coll *schemapb.CollectionSchema, schema *schemapb.FunctionSchema, extraInfo *models.ModelExtraInfo) (Runner, error) {
	switch schema.GetType() {
	case schemapb.FunctionType_BM25: // ignore bm25 function
		return nil, nil
	case schemapb.FunctionType_TextEmbedding:
		f, err := NewTextEmbeddingFunction(coll, schema, extraInfo)
		if err != nil {
			return nil, err
		}
		return f, nil
	case schemapb.FunctionType_MinHash:
		return nil, nil
	default:
		return nil, merr.WrapErrParameterInvalidMsg("unknown functionRunner type %s", schema.GetType().String())
	}
}

// validateFunction validates a single function.
func validateFunction(schema *schemapb.CollectionSchema, fSchema *schemapb.FunctionSchema, extraInfo *models.ModelExtraInfo) error {
	f, err := createFunction(schema, fSchema, extraInfo)
	if err != nil {
		return err
	}

	// BM25 or minhash function returns nil from createFunction
	if f == nil {
		if fSchema.GetType() == schemapb.FunctionType_MinHash {
			err := function.ValidateMinHashFunction(schema, fSchema)
			if err != nil {
				return err
			}
		}
		// bm25 or minhash validate pass
		return nil
	}

	if err := f.Check(context.Background()); err != nil {
		return merr.Wrapf(err, "check function [%s:%s] failed", fSchema.Name, fSchema.GetType().String())
	}
	return nil
}

// ValidateFunctions validates functions in the schema.
// If needValidateFunctionName is empty, all functions will be validated; otherwise, only the specified function will be validated.
func ValidateFunctions(schema *schemapb.CollectionSchema, needValidateFunctionName string, extraInfo *models.ModelExtraInfo) error {
	if needValidateFunctionName == "" {
		for _, fSchema := range schema.Functions {
			if err := validateFunction(schema, fSchema, extraInfo); err != nil {
				return err
			}
		}
		return nil
	}

	for _, fSchema := range schema.Functions {
		if fSchema.Name == needValidateFunctionName {
			return validateFunction(schema, fSchema, extraInfo)
		}
	}

	return merr.WrapErrParameterInvalidMsg("function [%s] not found in schema", needValidateFunctionName)
}

func NewFunctionExecutor(schema *schemapb.CollectionSchema, functions []*schemapb.FunctionSchema, extraInfo *models.ModelExtraInfo) (*FunctionExecutor, error) {
	executor := &FunctionExecutor{
		runners: make(map[int64]Runner),
	}
	if functions == nil {
		functions = schema.Functions
	}
	for _, fSchema := range functions {
		runner, err := createFunction(schema, fSchema, extraInfo)
		if err != nil {
			return nil, err
		}
		if runner != nil {
			executor.runners[fSchema.GetOutputFieldIds()[0]] = runner
		}
	}
	return executor, nil
}

func (executor *FunctionExecutor) processSingleFunction(ctx context.Context, runner Runner, msg *msgstream.InsertMsg) ([]*schemapb.FieldData, error) {
	inputs := make([]*schemapb.FieldData, 0, len(runner.GetSchema().GetInputFieldNames()))
	for _, name := range runner.GetSchema().GetInputFieldNames() {
		for _, field := range msg.FieldsData {
			if field.GetFieldName() == name {
				inputs = append(inputs, field)
			}
		}
	}
	if len(inputs) != len(runner.GetSchema().InputFieldIds) {
		return nil, merr.WrapErrParameterInvalidMsg("input field not found")
	}

	tr := timerecord.NewTimeRecorder("function ProcessInsert")
	outputs, err := runner.ProcessInsert(ctx, inputs)
	if err != nil {
		return nil, err
	}

	metrics.ProxyFunctionlatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), runner.GetCollectionName(), runner.GetFunctionTypeName(), runner.GetFunctionProvider(), runner.GetFunctionName()).Observe(float64(tr.RecordSpan().Milliseconds()))
	tr.CtxElapse(ctx, "function ProcessInsert done")
	return outputs, nil
}

func (executor *FunctionExecutor) ProcessInsert(ctx context.Context, msg *msgstream.InsertMsg) error {
	numRows := msg.NumRows
	for _, runner := range executor.runners {
		if numRows > uint64(runner.MaxBatch()) {
			return merr.WrapErrParameterInvalidMsg("numRows [%d] > function [%s]'s max batch [%d]", numRows, runner.GetSchema().Name, runner.MaxBatch())
		}
	}

	outputs := make(chan []*schemapb.FieldData, len(executor.runners))
	errChan := make(chan error, len(executor.runners))
	var wg sync.WaitGroup
	for _, runner := range executor.runners {
		wg.Add(1)
		go func(runner Runner) {
			defer wg.Done()
			data, err := executor.processSingleFunction(ctx, runner, msg)
			if err != nil {
				errChan <- err
				return
			}
			outputs <- data
		}(runner)
	}
	wg.Wait()
	close(errChan)
	close(outputs)

	// Collect all errors
	var errs []error
	for err := range errChan {
		errs = append(errs, err)
	}
	if len(errs) > 0 {
		return merr.WrapErrFunctionFailedMsg("function execution failed: %v", errs)
	}

	for output := range outputs {
		msg.FieldsData = append(msg.FieldsData, output...)
	}
	return nil
}

func (executor *FunctionExecutor) processSingleSearch(ctx context.Context, runner Runner, placeholderGroup []byte) ([]byte, error) {
	pb := &commonpb.PlaceholderGroup{}
	proto.Unmarshal(placeholderGroup, pb)
	if len(pb.Placeholders) != 1 {
		return nil, merr.WrapErrParameterInvalidMsg("No placeholders founded")
	}
	if pb.Placeholders[0].Type != commonpb.PlaceholderType_VarChar {
		return placeholderGroup, nil
	}

	tr := timerecord.NewTimeRecorder("function ProcessSearch")
	res, err := runner.ProcessSearch(ctx, pb)
	if err != nil {
		return nil, err
	}
	metrics.ProxyFunctionlatency.WithLabelValues(strconv.FormatInt(paramtable.GetNodeID(), 10), runner.GetCollectionName(), runner.GetFunctionTypeName(), runner.GetFunctionProvider(), runner.GetFunctionName()).Observe(float64(tr.RecordSpan().Milliseconds()))
	tr.CtxElapse(ctx, "function ProcessSearch done")
	return proto.Marshal(res)
}

func (executor *FunctionExecutor) prcessSearch(ctx context.Context, req *internalpb.SearchRequest) error {
	runner, exist := executor.runners[req.FieldId]
	if !exist {
		return merr.WrapErrParameterInvalidMsg("can not found function in field %d", req.FieldId)
	}
	if req.Nq > int64(runner.MaxBatch()) {
		return merr.WrapErrParameterInvalidMsg("nq [%d] > function [%s]'s max batch [%d]", req.Nq, runner.GetSchema().Name, runner.MaxBatch())
	}
	if newHolder, err := executor.processSingleSearch(ctx, runner, req.GetPlaceholderGroup()); err != nil {
		return err
	} else {
		req.PlaceholderGroup = newHolder
	}
	return nil
}

func (executor *FunctionExecutor) prcessAdvanceSearch(ctx context.Context, req *internalpb.SearchRequest) error {
	outputs := make(chan map[int64][]byte, len(req.GetSubReqs()))
	errChan := make(chan error, len(req.GetSubReqs()))
	var wg sync.WaitGroup
	for idx, sub := range req.GetSubReqs() {
		if runner, exist := executor.runners[sub.FieldId]; exist {
			if sub.Nq > int64(runner.MaxBatch()) {
				return merr.WrapErrParameterInvalidMsg("nq [%d] > function [%s]'s max batch [%d]", sub.Nq, runner.GetSchema().Name, runner.MaxBatch())
			}
			wg.Add(1)
			go func(runner Runner, idx int64, placeholderGroup []byte) {
				defer wg.Done()
				if newHolder, err := executor.processSingleSearch(ctx, runner, placeholderGroup); err != nil {
					errChan <- err
				} else {
					outputs <- map[int64][]byte{idx: newHolder}
				}
			}(runner, int64(idx), sub.GetPlaceholderGroup())
		}
	}
	wg.Wait()
	close(errChan)
	close(outputs)
	for err := range errChan {
		return err
	}

	for output := range outputs {
		for idx, holder := range output {
			req.SubReqs[idx].PlaceholderGroup = holder
		}
	}
	return nil
}

func (executor *FunctionExecutor) ProcessSearch(ctx context.Context, req *internalpb.SearchRequest) error {
	if !req.IsAdvanced {
		return executor.prcessSearch(ctx, req)
	}
	return executor.prcessAdvanceSearch(ctx, req)
}

func (executor *FunctionExecutor) processSingleBulkInsert(ctx context.Context, runner Runner, data *storage.InsertData) (map[storage.FieldID]storage.FieldData, error) {
	inputs := make([]storage.FieldData, 0, len(runner.GetSchema().InputFieldIds))
	for idx, id := range runner.GetSchema().InputFieldIds {
		field, exist := data.Data[id]
		if !exist {
			return nil, merr.WrapErrParameterInvalidMsg("can not find input field: [%s]", runner.GetSchema().GetInputFieldNames()[idx])
		}
		inputs = append(inputs, field)
	}

	outputs, err := runner.ProcessBulkInsert(ctx, inputs)
	if err != nil {
		return nil, err
	}
	return outputs, nil
}

func (executor *FunctionExecutor) ProcessBulkInsert(ctx context.Context, data *storage.InsertData) error {
	// Since concurrency has already been used in the outer layer, only a serial logic access model is used here.
	for _, runner := range executor.runners {
		output, err := executor.processSingleBulkInsert(ctx, runner, data)
		if err != nil {
			return err
		}
		for k, v := range output {
			data.Data[k] = v
		}
	}
	return nil
}