Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

509 lines
18 KiB
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 task
import (
"context"
"fmt"
"strconv"
"time"
"github.com/samber/lo"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v3/msgpb"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/utils"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/commonpbutil"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
// idSource helper type for using id as task source
type idSource int64
func (s idSource) String() string {
return fmt.Sprintf("ID-%d", s)
}
func WrapIDSource(id int64) Source {
return idSource(id)
}
func Wait(ctx context.Context, timeout time.Duration, tasks ...Task) error {
ctx, cancel := context.WithTimeout(ctx, timeout)
defer cancel()
var err error
go func() {
for _, task := range tasks {
err = task.Wait()
if err != nil {
cancel()
break
}
}
cancel()
}()
<-ctx.Done()
return err
}
func SetPriority(priority Priority, tasks ...Task) {
for i := range tasks {
tasks[i].SetPriority(priority)
}
}
func SetReason(reason string, tasks ...Task) {
for i := range tasks {
tasks[i].SetReason(reason)
}
}
// GetTaskType returns the task's type,
// for now, only 3 types;
// - only 1 grow action -> Grow
// - only 1 reduce action -> Reduce
// - 1 grow action, and ends with 1 reduce action -> Move
func GetTaskType(task Task) Type {
switch {
case len(task.Actions()) > 1:
return TaskTypeMove
case task.Actions()[0].Type() == ActionTypeGrow:
return TaskTypeGrow
case task.Actions()[0].Type() == ActionTypeReduce:
return TaskTypeReduce
case task.Actions()[0].Type() == ActionTypeUpdate || task.Actions()[0].Type() == ActionTypeReopen:
return TaskTypeUpdate
case task.Actions()[0].Type() == ActionTypeStatsUpdate:
return TaskTypeStatsUpdate
case task.Actions()[0].Type() == ActionTypeDropIndex:
return TaskTypeDropIndex
}
return 0
}
func mergeCollectionProps(schemaProps []*commonpb.KeyValuePair, collectionProps []*commonpb.KeyValuePair) []*commonpb.KeyValuePair {
// Merge the collectionProps and schemaProps maps, giving priority to the values in schemaProps if there are duplicate keys.
props := make(map[string]string)
for _, p := range collectionProps {
props[p.GetKey()] = p.GetValue()
}
for _, p := range schemaProps {
props[p.GetKey()] = p.GetValue()
}
var ret []*commonpb.KeyValuePair
for k, v := range props {
ret = append(ret, &commonpb.KeyValuePair{
Key: k,
Value: v,
})
}
return ret
}
func packLoadSegmentRequest(
task *SegmentTask,
action Action,
schema *schemapb.CollectionSchema,
collectionProperties []*commonpb.KeyValuePair,
loadMeta *querypb.LoadMetaInfo,
loadInfo *querypb.SegmentLoadInfo,
indexInfo []*indexpb.IndexInfo,
) *querypb.LoadSegmentsRequest {
loadScope := querypb.LoadScope_Full
if action.Type() == ActionTypeUpdate {
loadScope = querypb.LoadScope_Index
}
if action.Type() == ActionTypeStatsUpdate {
loadScope = querypb.LoadScope_Stats
}
if action.Type() == ActionTypeReopen {
loadScope = querypb.LoadScope_Reopen
}
if task.Source() == utils.LeaderChecker {
loadScope = querypb.LoadScope_Delta
}
finalSchema := applyCollectionSettings(schema, collectionProperties)
applyIndexWarmupSetting(loadInfo, finalSchema, collectionProperties)
return &querypb.LoadSegmentsRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_LoadSegments),
commonpbutil.WithMsgID(task.ID()),
),
Infos: []*querypb.SegmentLoadInfo{loadInfo},
Schema: finalSchema, // assign it for compatibility of rolling upgrade from 2.2.x to 2.3
LoadMeta: loadMeta, // assign it for compatibility of rolling upgrade from 2.2.x to 2.3
CollectionID: task.CollectionID(),
ReplicaID: task.ReplicaID(),
DeltaPositions: []*msgpb.MsgPosition{loadInfo.GetDeltaPosition()}, // assign it for compatibility of rolling upgrade from 2.2.x to 2.3
DstNodeID: action.Node(),
Version: time.Now().UnixNano(),
NeedTransfer: true,
IndexInfoList: indexInfo,
LoadScope: loadScope,
}
}
func packReleaseSegmentRequest(task *SegmentTask, action *SegmentAction) *querypb.ReleaseSegmentsRequest {
return &querypb.ReleaseSegmentsRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_ReleaseSegments),
commonpbutil.WithMsgID(task.ID()),
),
NodeID: action.Node(),
CollectionID: task.CollectionID(),
SegmentIDs: []int64{task.SegmentID()},
Scope: action.GetScope(),
Shard: action.GetShard(),
NeedTransfer: false,
}
}
func packLoadMeta(loadType querypb.LoadType, collectionInfo *milvuspb.DescribeCollectionResponse, resourceGroup string, loadFields []int64, partitions ...int64) *querypb.LoadMetaInfo {
return &querypb.LoadMetaInfo{
LoadType: loadType,
CollectionID: collectionInfo.GetCollectionID(),
PartitionIDs: partitions,
DbName: collectionInfo.GetDbName(),
ResourceGroup: resourceGroup,
LoadFields: loadFields,
// The update timestamp is a load barrier, not the logical schema version.
// QueryNode uses it only to reject stale load results after schema changes.
SchemaBarrierTs: collectionInfo.GetUpdateTimestamp(),
}
}
func packSubChannelRequest(
task *ChannelTask,
action Action,
schema *schemapb.CollectionSchema,
collectionProperties []*commonpb.KeyValuePair,
loadMeta *querypb.LoadMetaInfo,
channel *meta.DmChannel,
indexInfo []*indexpb.IndexInfo,
partitions []int64,
targetVersion int64,
) *querypb.WatchDmChannelsRequest {
finalSchema := applyCollectionSettings(schema, collectionProperties)
return &querypb.WatchDmChannelsRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_WatchDmChannels),
commonpbutil.WithMsgID(task.ID()),
),
NodeID: action.Node(),
CollectionID: task.CollectionID(),
PartitionIDs: partitions,
Infos: []*datapb.VchannelInfo{channel.VchannelInfo},
Schema: finalSchema, // assign it for compatibility of rolling upgrade from 2.2.x to 2.3
LoadMeta: loadMeta, // assign it for compatibility of rolling upgrade from 2.2.x to 2.3
ReplicaID: task.ReplicaID(),
Version: time.Now().UnixNano(),
IndexInfoList: indexInfo,
TargetVersion: targetVersion,
}
}
func fillSubChannelRequest(
ctx context.Context,
req *querypb.WatchDmChannelsRequest,
broker meta.Broker,
includeFlushed bool,
) error {
segmentIDs := typeutil.NewUniqueSet()
for _, vchannel := range req.GetInfos() {
if includeFlushed {
segmentIDs.Insert(vchannel.GetFlushedSegmentIds()...)
}
segmentIDs.Insert(vchannel.GetUnflushedSegmentIds()...)
segmentIDs.Insert(vchannel.GetLevelZeroSegmentIds()...)
}
if segmentIDs.Len() == 0 {
return nil
}
segmentInfos, err := broker.GetSegmentInfo(ctx, segmentIDs.Collect()...)
if err != nil {
return err
}
req.SegmentInfos = lo.SliceToMap(segmentInfos, func(info *datapb.SegmentInfo) (int64, *datapb.SegmentInfo) {
return info.GetID(), info
})
return nil
}
func packUnsubDmChannelRequest(task *ChannelTask, action Action) *querypb.UnsubDmChannelRequest {
return &querypb.UnsubDmChannelRequest{
Base: commonpbutil.NewMsgBase(
commonpbutil.WithMsgType(commonpb.MsgType_UnsubDmChannel),
commonpbutil.WithMsgID(task.ID()),
),
NodeID: action.Node(),
CollectionID: task.CollectionID(),
ChannelName: task.Channel(),
}
}
// applyCollectionSettings applies collection-level mmap and warmup settings to all fields.
// It combines applyCollectionMmapSetting and applyCollectionWarmupSetting into a single call.
func applyCollectionSettings(schema *schemapb.CollectionSchema,
collectionProperties []*commonpb.KeyValuePair,
) *schemapb.CollectionSchema {
// first clone the schema then apply some settings to it
schemaCloned := typeutil.Clone(schema)
schemaCloned.Properties = mergeCollectionProps(schemaCloned.Properties, collectionProperties)
schemaCloned = applyCollectionMmapSetting(schemaCloned, collectionProperties)
schemaCloned = applyCollectionWarmupSetting(schemaCloned, collectionProperties)
return schemaCloned
}
func applyCollectionMmapSetting(schema *schemapb.CollectionSchema,
collectionProperties []*commonpb.KeyValuePair,
) *schemapb.CollectionSchema {
// field mmap enabled if collection-level mmap enabled or the field mmap enabled
collectionMmapEnabled, exist := common.IsMmapDataEnabled(collectionProperties...)
for _, field := range schema.GetFields() {
if exist &&
// field-level mmap setting has higher priority than collection-level mmap setting, skip if field-level mmap enabled
!common.FieldHasMmapKey(schema, field.GetFieldID()) {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.MmapEnabledKey,
Value: strconv.FormatBool(collectionMmapEnabled),
})
}
}
for _, structField := range schema.GetStructArrayFields() {
structTypeParams := structField.GetTypeParams()
structMmapEnabled, structExist := common.IsMmapDataEnabled(structTypeParams...)
// If struct field itself doesn't have mmap setting, inherit from collection
if !structExist && exist &&
!common.FieldHasMmapKey(schema, structField.GetFieldID()) {
structField.TypeParams = append(structField.TypeParams, &commonpb.KeyValuePair{
Key: common.MmapEnabledKey,
Value: strconv.FormatBool(collectionMmapEnabled),
})
// Update struct's mmap state since we just set it
structMmapEnabled = collectionMmapEnabled
structExist = true
}
// Apply mmap setting to fields inside struct
for _, field := range structField.GetFields() {
// Skip if field already has mmap setting
if common.FieldHasMmapKey(schema, field.GetFieldID()) {
continue
}
// Priority: struct field setting > collection setting
if structExist {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.MmapEnabledKey,
Value: strconv.FormatBool(structMmapEnabled),
})
} else if exist {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.MmapEnabledKey,
Value: strconv.FormatBool(collectionMmapEnabled),
})
}
}
}
return schema
}
// autoWarmupForNonPKIsolationCollection checks if autoWarmupForNonPKIsolationCollection should be applied for a collection.
// Returns true when the AutoWarmupForNonPKIsolationCollection config is enabled AND the collection
// does NOT have partition key isolation.
func autoWarmupForNonPKIsolationCollection(collectionProperties []*commonpb.KeyValuePair) bool {
if !paramtable.Get().QueryCoordCfg.AutoWarmupForNonPKIsolationCollection.GetAsBool() {
return false
}
isPKI, isError := common.IsPartitionKeyIsolationKvEnabled(collectionProperties...)
if isError != nil {
mlog.Warn(context.TODO(), "failed to parse partition key isolation, autowarmup is disabled", mlog.Err(isError))
return false
}
if !isPKI {
mlog.Info(context.TODO(), "collection is not partition key isolated and autowarmup is enabled, force scalar field/index and vector index warmup to sync")
}
return !isPKI
}
// applyCollectionWarmupSetting applies collection-level warmup setting to all fields
// and propagates struct-level warmup to nested fields.
// Priority: field-level > struct-level > collection-level > autoWarmupForNonPKIsolationCollection (scalar only)
// Collection-level granular keys: warmup.scalarField, warmup.vectorField
func applyCollectionWarmupSetting(schema *schemapb.CollectionSchema,
collectionProperties []*commonpb.KeyValuePair,
) *schemapb.CollectionSchema {
// Get collection-level granular warmup policies
scalarFieldWarmup, scalarFieldExist := common.GetWarmupPolicyByKey(common.WarmupScalarFieldKey, collectionProperties...)
vectorFieldWarmup, vectorFieldExist := common.GetWarmupPolicyByKey(common.WarmupVectorFieldKey, collectionProperties...)
autoWarmup := autoWarmupForNonPKIsolationCollection(collectionProperties)
// Apply collection-level warmup to regular fields
for _, field := range schema.GetFields() {
// field-level warmup setting has higher priority, skip if field already has warmup setting
if common.FieldHasWarmupKey(schema, field.GetFieldID()) {
continue
}
isVector := typeutil.IsVectorType(field.GetDataType())
if isVector && vectorFieldExist {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: vectorFieldWarmup,
})
} else if !isVector && scalarFieldExist {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: scalarFieldWarmup,
})
} else if autoWarmup && !isVector {
// autoWarmupForNonPKIsolationCollection fallback: force sync warmup for scalar fields (not vector fields)
// vector fields are excluded from autowarmup due to its large size
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: common.WarmupSync,
})
}
}
// Apply warmup to struct array fields and their nested fields
for _, structField := range schema.GetStructArrayFields() {
// Get struct field's warmup setting
structFieldWarmup, structHasWarmup := common.GetWarmupPolicy(structField.GetTypeParams()...)
// Apply warmup setting to fields inside struct
for _, field := range structField.GetFields() {
// Skip if field already has warmup setting
if common.FieldHasWarmupKey(schema, field.GetFieldID()) {
continue
}
// Priority: struct field setting > collection setting > autoWarmupForNonPKIsolationCollection
if structHasWarmup {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: structFieldWarmup,
})
} else {
isVector := typeutil.IsVectorType(field.GetDataType())
if isVector && vectorFieldExist {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: vectorFieldWarmup,
})
} else if !isVector && scalarFieldExist {
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: scalarFieldWarmup,
})
} else if autoWarmup && !isVector {
// autoWarmupForNonPKIsolationCollection fallback for struct nested scalar fields
// vector fields are excluded from autowarmup due to its large size
field.TypeParams = append(field.TypeParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: common.WarmupSync,
})
}
}
}
}
return schema
}
// applyIndexWarmupSetting applies collection-level index warmup setting to segment index params
// Index params warmup setting has higher priority than collection-level
// Priority: index-level > collection-level > autoWarmupForNonPKIsolationCollection (all indexes)
// Collection-level granular keys: warmup.scalarIndex, warmup.vectorIndex
func applyIndexWarmupSetting(loadInfo *querypb.SegmentLoadInfo, schema *schemapb.CollectionSchema, collectionProperties []*commonpb.KeyValuePair) {
// Get collection-level granular warmup policies for indexes
scalarIndexWarmup, scalarIndexExist := common.GetWarmupPolicyByKey(common.WarmupScalarIndexKey, collectionProperties...)
vectorIndexWarmup, vectorIndexExist := common.GetWarmupPolicyByKey(common.WarmupVectorIndexKey, collectionProperties...)
autoWarmup := autoWarmupForNonPKIsolationCollection(collectionProperties)
if !scalarIndexExist && !vectorIndexExist && !autoWarmup {
return
}
// Build fieldID to field schema map for quick lookup
fieldMap := make(map[int64]*schemapb.FieldSchema)
for _, field := range schema.GetFields() {
fieldMap[field.GetFieldID()] = field
}
// Include nested fields from struct arrays (struct array fields themselves don't support indexes)
for _, structField := range schema.GetStructArrayFields() {
for _, field := range structField.GetFields() {
fieldMap[field.GetFieldID()] = field
}
}
for _, indexInfo := range loadInfo.GetIndexInfos() {
// Check if index params already has warmup setting
_, exist := common.GetWarmupPolicy(indexInfo.IndexParams...)
if exist {
continue
}
field, ok := fieldMap[indexInfo.GetFieldID()]
if !ok {
continue
}
isVector := typeutil.IsVectorType(field.GetDataType())
if isVector && vectorIndexExist {
indexInfo.IndexParams = append(indexInfo.IndexParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: vectorIndexWarmup,
})
} else if !isVector && scalarIndexExist {
indexInfo.IndexParams = append(indexInfo.IndexParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: scalarIndexWarmup,
})
} else if autoWarmup {
// autoWarmupForNonPKIsolationCollection fallback: force sync warmup for ALL indexes (scalar and vector)
// here vector indexes are included in autowarmup bcz they are critical for search performance
indexInfo.IndexParams = append(indexInfo.IndexParams, &commonpb.KeyValuePair{
Key: common.WarmupKey,
Value: common.WarmupSync,
})
}
}
}