furyhawk/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-07-21 10:15:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

enhance: Reduce segment load resource lock contention (#50340)

Browse Source 
Fixes #50334

## What changed
- Move segment loading resource estimation outside the committed
resource lock to reduce lock hold time.
- Re-sample physical memory and disk usage after estimation while
holding the lock, so physical usage and committed loading resource are
checked consistently.
- Fix batched GPU memory accounting and release tiered eviction loading
resource if GPU validation fails after reservation.

## Verification
- gofumpt -l internal/querynodev2/segments/segment_loader.go
internal/querynodev2/segments/segment_loader_test.go
- git diff --check 1cc43cdef8303a92584c8e84afdd15fdabdef8eb..HEAD
- go test -v -count=1 -tags dynamic,test -gcflags="all=-N -l"
-ldflags="-r ${MILVUS_WORK_DIR}/cmake_build/lib -r
${MILVUS_WORK_DIR}/internal/core/output/lib"
./internal/querynodev2/segments -run
"TestSegmentLoader/TestSegmentLoaderDetailSuite/(TestCheckLoadingResourceWithDiskLimit|TestCheckLoadingResourceWithMemoryLimit|TestCheckSegmentGpuMemSizeWithBatchedEstimates|TestRequestResource)"
-timeout 300s (blocked locally: Package milvus_core was not found in
pkg-config search path)

---------

Signed-off-by: sijie-ni-0214 <sijie.ni@zilliz.com>
This commit is contained in:
sijie-ni-0214
2026-07-03 06:26:29 +08:00
committed by GitHub
parent c62963c6ad
commit 6b710e80c4
2 changed files with 252 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/querynodev2/segments/segment_loader.go
+79 -39
View File
@@ -515,6 +515,22 @@ func (loader *segmentLoader) requestResource(ctx context.Context, infos ...*quer
return requestResourceResult{}, nil
}
segmentIDs := lo.Map(infos, func(info *querypb.SegmentLoadInfo, _ int) int64 {
return info.GetSegmentID()
})
logger := mlog.With(
mlog.Int64s("segmentIDs", segmentIDs),
)
loadingUsage, maxSegmentSize, err := loader.estimateSegmentLoadingResourceUsage(ctx, infos...)
if err != nil {
logger.Warn(ctx, "no sufficient physical resource to load segments", mlog.Err(err))
return requestResourceResult{}, err
}
loader.mut.Lock()
defer loader.mut.Unlock()
physicalMemoryUsage := hardware.GetUsedMemoryCount()
totalMemory := hardware.GetMemoryCount()
@@ -524,9 +540,6 @@ func (loader *segmentLoader) requestResource(ctx context.Context, infos ...*quer
}
diskCap := paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsUint64()
loader.mut.Lock()
defer loader.mut.Unlock()
result := requestResourceResult{
CommittedResource: loader.committedResource,
}
@@ -546,15 +559,12 @@ func (loader *segmentLoader) requestResource(ctx context.Context, infos ...*quer
// return result, err
// }
// then get physical resource usage for loading segments
mu, du, err := loader.checkSegmentSize(ctx, infos, totalMemory, physicalMemoryUsage, physicalDiskUsage)
if err != nil {
mlog.Warn(context.TODO(), "no sufficient physical resource to load segments", mlog.Err(err))
if err := loader.checkLoadingResource(ctx, logger, loadingUsage, maxSegmentSize, totalMemory, physicalMemoryUsage, physicalDiskUsage); err != nil {
return result, err
}
result.Resource.MemorySize = mu
result.Resource.DiskSize = du
result.Resource.MemorySize = loadingUsage.MemorySize
result.Resource.DiskSize = loadingUsage.DiskSize
// result.LogicalResource.MemorySize = lmu
// result.LogicalResource.DiskSize = ldu
@@ -1806,20 +1816,14 @@ func (loader *segmentLoader) checkLogicalSegmentSize(ctx context.Context, segmen
return predictLogicalMemUsage - logicalMemUsage, predictLogicalDiskUsage - logicalDiskUsage, nil
}
// checkSegmentSize checks whether the memory & disk is sufficient to load the segments
// returns the memory & disk usage while loading if possible to load,
// otherwise, returns error
func (loader *segmentLoader) checkSegmentSize(ctx context.Context, segmentLoadInfos []*querypb.SegmentLoadInfo, totalMem, memUsage uint64, localDiskUsage int64) (uint64, uint64, error) {
func (loader *segmentLoader) estimateSegmentLoadingResourceUsage(ctx context.Context, segmentLoadInfos ...*querypb.SegmentLoadInfo) (*ResourceUsage, uint64, error) {
if len(segmentLoadInfos) == 0 {
return 0, 0, nil
return &ResourceUsage{}, 0, nil
}
memUsage = memUsage + loader.committedResource.MemorySize
if memUsage == 0 || totalMem == 0 {
return 0, 0, merr.WrapErrServiceInternalMsg("get memory failed when checkSegmentSize")
}
diskUsage := uint64(localDiskUsage) + loader.committedResource.DiskSize
logger := mlog.With(
mlog.Int64("collectionID", segmentLoadInfos[0].GetCollectionID()),
)
maxFactor := resourceEstimateFactor{
memoryUsageFactor: paramtable.Get().QueryNodeCfg.LoadMemoryUsageFactor.GetAsFloat(),
@@ -1833,22 +1837,22 @@ func (loader *segmentLoader) checkSegmentSize(ctx context.Context, segmentLoadIn
externalRawDataFactor: paramtable.Get().QueryNodeCfg.ExternalCollectionRawDataFactor.GetAsFloat(),
}
maxSegmentSize := uint64(0)
predictMemUsage := memUsage
predictDiskUsage := diskUsage
predictMemUsage := uint64(0)
predictDiskUsage := uint64(0)
var predictGpuMemUsage []uint64
mmapFieldCount := 0
for _, loadInfo := range segmentLoadInfos {
collection := loader.manager.Collection.Get(loadInfo.GetCollectionID())
loadingUsage, err := estimateLoadingResourceUsageOfSegment(collection.Schema(), loadInfo, maxFactor)
if err != nil {
mlog.Warn(context.TODO(), "failed to estimate max resource usage of segment",
logger.Warn(ctx, "failed to estimate max resource usage of segment",
mlog.Int64("collectionID", loadInfo.GetCollectionID()),
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
mlog.Err(err))
return 0, 0, err
return nil, 0, err
}
mlog.Debug(context.TODO(), "segment resource for loading",
logger.Debug(ctx, "segment resource for loading",
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
mlog.Float64("loadingMemoryUsage(MB)", logutil.ToMB(float64(loadingUsage.MemorySize))),
mlog.Float64("loadingDiskUsage(MB)", logutil.ToMB(float64(loadingUsage.DiskSize))),
@@ -1857,13 +1861,41 @@ func (loader *segmentLoader) checkSegmentSize(ctx context.Context, segmentLoadIn
mmapFieldCount += loadingUsage.MmapFieldCount
predictDiskUsage += loadingUsage.DiskSize
predictMemUsage += loadingUsage.MemorySize
predictGpuMemUsage = loadingUsage.FieldGpuMemorySize
predictGpuMemUsage = append(predictGpuMemUsage, loadingUsage.FieldGpuMemorySize...)
if loadingUsage.MemorySize > maxSegmentSize {
maxSegmentSize = loadingUsage.MemorySize
}
}
mlog.Info(context.TODO(), "predict memory and disk usage while loading (in MiB)",
return &ResourceUsage{
MemorySize: predictMemUsage,
DiskSize: predictDiskUsage,
MmapFieldCount: mmapFieldCount,
FieldGpuMemorySize: predictGpuMemUsage,
}, maxSegmentSize, nil
}
// checkLoadingResource checks physical resource limits for an already-estimated loading usage.
// Callers that race with load resource commits must hold loader.mut.
func (loader *segmentLoader) checkLoadingResource(
ctx context.Context,
logger *mlog.Logger,
loadingUsage *ResourceUsage,
maxSegmentSize uint64,
totalMem uint64,
memUsage uint64,
localDiskUsage int64,
) error {
memUsage += loader.committedResource.MemorySize
if memUsage == 0 || totalMem == 0 {
return merr.WrapErrServiceInternalMsg("get memory failed when checkLoadingResource")
}
diskUsage := uint64(localDiskUsage) + loader.committedResource.DiskSize
predictMemUsage := memUsage + loadingUsage.MemorySize
predictDiskUsage := diskUsage + loadingUsage.DiskSize
logger.Info(ctx, "predict memory and disk usage while loading (in MiB)",
mlog.Float64("maxSegmentSize(MB)", logutil.ToMB(float64(maxSegmentSize))),
mlog.Float64("committedMemSize(MB)", logutil.ToMB(float64(loader.committedResource.MemorySize))),
mlog.Float64("memLimit(MB)", logutil.ToMB(float64(totalMem))),
@@ -1872,16 +1904,20 @@ func (loader *segmentLoader) checkSegmentSize(ctx context.Context, segmentLoadIn
mlog.Float64("diskUsage(MB)", logutil.ToMB(float64(diskUsage))),
mlog.Float64("predictMemUsage(MB)", logutil.ToMB(float64(predictMemUsage))),
mlog.Float64("predictDiskUsage(MB)", logutil.ToMB(float64(predictDiskUsage))),
mlog.Int("mmapFieldCount", mmapFieldCount),
mlog.Int("mmapFieldCount", loadingUsage.MmapFieldCount),
)
var loadingResource C.CResourceUsage
reservedLoadingResource := false
if paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool() {
loadingResource = C.CResourceUsage{
memory_bytes: C.int64_t(loadingUsage.MemorySize),
disk_bytes: C.int64_t(loadingUsage.DiskSize),
}
// try to reserve loading resource from caching layer
if ok := C.TryReserveLoadingResourceWithTimeout(C.CResourceUsage{
memory_bytes: C.int64_t(predictMemUsage - memUsage),
disk_bytes: C.int64_t(predictDiskUsage - diskUsage),
}, 1000); !ok {
return 0, 0, merr.WrapErrSegmentRequestResourceFailed("memory/disk",
if ok := C.TryReserveLoadingResourceWithTimeout(loadingResource, 1000); !ok {
return merr.WrapErrSegmentRequestResourceFailed("memory/disk",
fmt.Sprintf("failed to reserve loading resource from caching layer, predictMemUsage = %v MB, predictDiskUsage = %v MB, memUsage = %v MB, diskUsage = %v MB, memoryThresholdFactor = %f, diskThresholdFactor = %f",
logutil.ToMB(float64(predictMemUsage)),
logutil.ToMB(float64(predictDiskUsage)),
@@ -1891,6 +1927,7 @@ func (loader *segmentLoader) checkSegmentSize(ctx context.Context, segmentLoadIn
paramtable.Get().QueryNodeCfg.MaxDiskUsagePercentage.GetAsFloat(),
))
}
reservedLoadingResource = true
} else {
// fallback to original segment loading logic
if predictMemUsage > uint64(float64(totalMem)*paramtable.Get().QueryNodeCfg.OverloadedMemoryThresholdPercentage.GetAsFloat()) {
@@ -1902,7 +1939,7 @@ func (loader *segmentLoader) checkSegmentSize(ctx context.Context, segmentLoadIn
mlog.Float64("totalMemMB", logutil.ToMB(float64(totalMem))),
mlog.Float64("thresholdFactor", paramtable.Get().QueryNodeCfg.OverloadedMemoryThresholdPercentage.GetAsFloat()),
)
return 0, 0, merr.WrapErrSegmentRequestResourceFailed("Memory")
return merr.WrapErrSegmentRequestResourceFailed("Memory")
}
if predictDiskUsage > uint64(float64(paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsInt64())*paramtable.Get().QueryNodeCfg.MaxDiskUsagePercentage.GetAsFloat()) {
@@ -1913,16 +1950,19 @@ func (loader *segmentLoader) checkSegmentSize(ctx context.Context, segmentLoadIn
mlog.Float64("totalDiskMB", logutil.ToMB(float64(uint64(paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsInt64())))),
mlog.Float64("thresholdFactor", paramtable.Get().QueryNodeCfg.MaxDiskUsagePercentage.GetAsFloat()),
)
return 0, 0, merr.WrapErrSegmentRequestResourceFailed("Disk")
return merr.WrapErrSegmentRequestResourceFailed("Disk")
}
}
err := checkSegmentGpuMemSize(predictGpuMemUsage, float32(paramtable.Get().GpuConfig.OverloadedMemoryThresholdPercentage.GetAsFloat()))
err := checkSegmentGpuMemSize(loadingUsage.FieldGpuMemorySize, float32(paramtable.Get().GpuConfig.OverloadedMemoryThresholdPercentage.GetAsFloat()))
if err != nil {
return 0, 0, err
if reservedLoadingResource {
C.ReleaseLoadingResource(loadingResource)
}
return err
}
return predictMemUsage - memUsage, predictDiskUsage - diskUsage, nil
return nil
}
// this function is used to estimate the logical resource usage of a segment, which should only be used when tiered eviction is enabled
@@ -2637,7 +2677,7 @@ func checkSegmentGpuMemSize(fieldGpuMemSizeList []uint64, OverloadedMemoryThresh
)
return merr.WrapErrSegmentRequestResourceFailed("GPU")
}
currentGpuMem[minId] += minGpuMem
currentGpuMem[minId] = minGpuMem
}
return nil
}
internal/querynodev2/segments/segment_loader_test.go
+173 -4
View File
@@ -38,10 +38,12 @@ import (
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
"github.com/milvus-io/milvus/internal/util/initcore"
"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/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/hardware"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/metric"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -1733,9 +1735,159 @@ func (suite *SegmentLoaderDetailSuite) TestRequestResource() {
suite.NoError(err)
suite.EqualValues(1100000, resource.Resource.MemorySize)
})
suite.Run("commits_estimated_resource", func() {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.Key, "2")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.Key)
patchUsedMemory := mockey.Mock(hardware.GetUsedMemoryCount).To(func() uint64 {
return 100
}).Build()
defer patchUsedMemory.UnPatch()
patchTotalMemory := mockey.Mock(hardware.GetMemoryCount).To(func() uint64 {
return 1024 * 1024 * 1024
}).Build()
defer patchTotalMemory.UnPatch()
patchCPU := mockey.Mock(hardware.GetCPUNum).To(func() int {
return 8
}).Build()
defer patchCPU.UnPatch()
suite.loader.duf.usage.Store(200)
suite.loader.committedResource = LoadResource{
MemorySize: 10,
DiskSize: 20,
}
resource, err := suite.loader.requestResource(context.Background(), loadInfo)
suite.NoError(err)
suite.EqualValues(44000, resource.Resource.MemorySize)
suite.EqualValues(10, resource.CommittedResource.MemorySize)
suite.EqualValues(20, resource.CommittedResource.DiskSize)
suite.Equal(1, resource.ConcurrencyLevel)
suite.EqualValues(44010, suite.loader.committedResource.MemorySize)
suite.EqualValues(20, suite.loader.committedResource.DiskSize)
})
suite.Run("estimate_failed", func() {
invalidCollectionID := suite.collectionID + 1
invalidSchema := &schemapb.CollectionSchema{
Name: "invalid",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 100, Name: "duplicated", DataType: schemapb.DataType_Int64},
},
}
originalCollectionManager := suite.manager.Collection
collectionManager := NewMockCollectionManager(suite.T())
collectionManager.EXPECT().Get(invalidCollectionID).
Return(NewCollectionWithoutSegcoreForTest(invalidCollectionID, invalidSchema))
suite.manager.Collection = collectionManager
defer func() {
suite.manager.Collection = originalCollectionManager
}()
_, err := suite.loader.requestResource(context.Background(), &querypb.SegmentLoadInfo{
SegmentID: 101,
CollectionID: invalidCollectionID,
})
suite.Error(err)
})
}
func (suite *SegmentLoaderDetailSuite) TestCheckSegmentSizeWithDiskLimit() {
func (suite *SegmentLoaderDetailSuite) TestEstimateSegmentLoadingResourceUsage() {
ctx := context.Background()
emptyUsage, maxSegmentSize, err := suite.loader.estimateSegmentLoadingResourceUsage(ctx)
suite.NoError(err)
suite.Zero(emptyUsage.MemorySize)
suite.Zero(emptyUsage.DiskSize)
suite.Zero(maxSegmentSize)
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.Key, "2")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.Key)
loadInfo1 := &querypb.SegmentLoadInfo{
SegmentID: 100,
CollectionID: suite.collectionID,
Level: datapb.SegmentLevel_L0,
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{LogSize: 10, MemorySize: 10},
},
},
},
}
loadInfo2 := &querypb.SegmentLoadInfo{
SegmentID: 101,
CollectionID: suite.collectionID,
Level: datapb.SegmentLevel_L0,
Deltalogs: []*datapb.FieldBinlog{
{
Binlogs: []*datapb.Binlog{
{LogSize: 30, MemorySize: 30},
},
},
},
}
usage, maxSegmentSize, err := suite.loader.estimateSegmentLoadingResourceUsage(ctx, loadInfo1, loadInfo2)
suite.NoError(err)
suite.EqualValues(80, usage.MemorySize)
suite.Zero(usage.DiskSize)
suite.Zero(usage.MmapFieldCount)
suite.Empty(usage.FieldGpuMemorySize)
suite.EqualValues(60, maxSegmentSize)
}
func (suite *SegmentLoaderDetailSuite) TestCheckLoadingResourceWithGpuLimit() {
ctx := context.Background()
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.Key)
patchGpu := mockey.Mock(hardware.GetAllGPUMemoryInfo).To(func() ([]hardware.GPUMemoryInfo, error) {
return []hardware.GPUMemoryInfo{{TotalMemory: 100, FreeMemory: 100}}, nil
}).Build()
defer patchGpu.UnPatch()
err := suite.loader.checkLoadingResource(ctx, mlog.With(), &ResourceUsage{
MemorySize: 10,
FieldGpuMemorySize: []uint64{100},
}, 10, 1000, 100, 0)
suite.Error(err)
suite.True(errors.Is(err, merr.ErrSegmentRequestResourceFailed))
}
func (suite *SegmentLoaderDetailSuite) TestCheckLoadingResourceReleasesTieredReservationOnGpuLimit() {
ctx := context.Background()
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.Key, "true")
suite.Require().NoError(initcore.InitTieredStorage(paramtable.Get()))
defer func() {
paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.Key)
suite.Require().NoError(initcore.InitTieredStorage(paramtable.Get()))
}()
patchGpu := mockey.Mock(hardware.GetAllGPUMemoryInfo).To(func() ([]hardware.GPUMemoryInfo, error) {
return []hardware.GPUMemoryInfo{{TotalMemory: 100, FreeMemory: 100}}, nil
}).Build()
defer patchGpu.UnPatch()
err := suite.loader.checkLoadingResource(ctx, mlog.With(), &ResourceUsage{
MemorySize: 10,
DiskSize: 10,
FieldGpuMemorySize: []uint64{100},
}, 10, 1000, 100, 0)
suite.Error(err)
suite.True(errors.Is(err, merr.ErrSegmentRequestResourceFailed))
}
func (suite *SegmentLoaderDetailSuite) TestCheckLoadingResourceWithDiskLimit() {
ctx := context.Background()
// Save original value and restore after test
@@ -1784,12 +1936,15 @@ func (suite *SegmentLoaderDetailSuite) TestCheckSegmentSizeWithDiskLimit() {
totalMem := uint64(1024 * 1024 * 1024) // 1GB
localDiskUsage := int64(100 * 1024) // 100KB
_, _, err := suite.loader.checkSegmentSize(ctx, []*querypb.SegmentLoadInfo{loadInfo}, memUsage, totalMem, localDiskUsage)
loadingUsage, maxSegmentSize, err := suite.loader.estimateSegmentLoadingResourceUsage(ctx, loadInfo)
suite.NoError(err)
err = suite.loader.checkLoadingResource(ctx, mlog.With(), loadingUsage, maxSegmentSize, totalMem, memUsage, localDiskUsage)
suite.Error(err)
suite.True(errors.Is(err, merr.ErrSegmentRequestResourceFailed))
}
func (suite *SegmentLoaderDetailSuite) TestCheckSegmentSizeWithMemoryLimit() {
func (suite *SegmentLoaderDetailSuite) TestCheckLoadingResourceWithMemoryLimit() {
ctx := context.Background()
// Create a test segment that would exceed the memory limit
@@ -1818,12 +1973,26 @@ func (suite *SegmentLoaderDetailSuite) TestCheckSegmentSizeWithMemoryLimit() {
// Set memory threshold to 80%
paramtable.Get().Save("queryNode.overloadedMemoryThresholdPercentage", "0.8")
defer paramtable.Get().Reset("queryNode.overloadedMemoryThresholdPercentage")
_, _, err := suite.loader.checkSegmentSize(ctx, []*querypb.SegmentLoadInfo{loadInfo}, memUsage, totalMem, localDiskUsage)
loadingUsage, maxSegmentSize, err := suite.loader.estimateSegmentLoadingResourceUsage(ctx, loadInfo)
suite.NoError(err)
err = suite.loader.checkLoadingResource(ctx, mlog.With(), loadingUsage, maxSegmentSize, totalMem, memUsage, localDiskUsage)
suite.Error(err)
suite.True(errors.Is(err, merr.ErrSegmentRequestResourceFailed))
}
func (suite *SegmentLoaderDetailSuite) TestCheckSegmentGpuMemSizeWithBatchedEstimates() {
patch := mockey.Mock(hardware.GetAllGPUMemoryInfo).To(func() ([]hardware.GPUMemoryInfo, error) {
return []hardware.GPUMemoryInfo{{TotalMemory: 100, FreeMemory: 100}}, nil
}).Build()
defer patch.UnPatch()
err := checkSegmentGpuMemSize([]uint64{30, 30, 30}, 1.0)
suite.NoError(err)
}
// SegmentLoaderTextIndexEstimateSuite tests resource estimation for text index (TextStatsLogs).
// These tests directly call estimateLoadingResourceUsageOfSegment and
// estimateLogicalResourceUsageOfSegment to verify that TextStatsLogs are correctly