Skip to content

Commit 4681c4c

Browse files
sparknacksparknack
committed
enhance: add storage v3 async load switch
Signed-off-by: Shawn Wang <shawn.wang@zilliz.com>
1 parent ead9492 commit 4681c4c

11 files changed

Lines changed: 144 additions & 17 deletions

File tree

‎configs/milvus.yaml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -578,6 +578,7 @@ queryNode:
578578
knowhereScoreConsistency: false # Enable knowhere strong consistency score computation logic
579579
storageV2:
580580
cellTargetSizeBytes: 4194304 # Target average byte size per storage v2 cache cell. Parquet row groups are greedily packed so that rgs_per_cell * avg_row_group_size ≈ this target. Each cell always contains at least one row group and cells never cross file boundaries. Tune larger for bigger batch IO (fewer cells) or smaller to get finer cache granularity. Default 4 MiB.
581+
enableAsyncLoad: true # Enable the storage v3 async column-group load translator. Disable to compare with the original synchronous translator.
581582
deleteDumpBatchSize: 10000 # Batch size for delete snapshot dump in segcore.
582583
loadMemoryUsageFactor: 1 # The multiply factor of calculating the memory usage while loading segments
583584
enableDisk: false # enable querynode load disk index, and search on disk index

‎internal/core/src/common/init_c.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
#include "log/Log.h"
3535
#include "segcore/memory_planner.h"
3636
#include "segcore/storagev2translator/GroupCTMeta.h"
37+
#include "segcore/storagev2translator/StorageV2Config.h"
3738
#include "storage/EntryStreamUtils.h"
3839
#include "storage/ThreadPool.h"
3940

@@ -217,6 +218,11 @@ SetStorageV2CellTargetSizeBytes(int64_t bytes) {
217218
milvus::segcore::storagev2translator::SetCellTargetSizeBytes(bytes);
218219
}
219220

221+
void
222+
SetStorageV2AsyncLoadEnabled(bool enabled) {
223+
milvus::segcore::storagev2translator::SetStorageV2AsyncLoadEnabled(enabled);
224+
}
225+
220226
void
221227
LogOpenSSLFIPSStatus() {
222228
std::call_once(fipsFlag, []() {

‎internal/core/src/common/init_c.h‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,11 @@ UpdateArrowIOThreadPoolMetrics();
120120
void
121121
SetStorageV2CellTargetSizeBytes(int64_t bytes);
122122

123+
// Enable the storage v3 async column-group load translator. Disable to compare
124+
// with the original synchronous translator.
125+
void
126+
SetStorageV2AsyncLoadEnabled(bool enabled);
127+
123128
#ifdef __cplusplus
124129
};
125130
#endif

‎internal/core/src/segcore/ChunkedSegmentSealedImpl.cpp‎

Lines changed: 39 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,7 @@
134134
#include "segcore/storagev2translator/GroupChunkTranslator.h"
135135
#include "segcore/storagev2translator/ManifestGroupTranslator.h"
136136
#include "segcore/storagev2translator/ManifestGroupTranslatorV2.h"
137+
#include "segcore/storagev2translator/StorageV2Config.h"
137138
#include "segcore/TextColumnCache.h"
138139
#include "storage/FileManager.h"
139140
#include "storage/KeyRetriever.h"
@@ -5469,23 +5470,44 @@ ChunkedSegmentSealedImpl::LoadColumnGroup(
54695470
cache_key_suffix = std::to_string(milvus_field_ids.front().get());
54705471
}
54715472

5472-
auto translator =
5473-
std::make_unique<storagev2translator::ManifestGroupTranslatorV2>(
5474-
get_segment_id(),
5475-
GroupChunkType::DEFAULT,
5476-
index,
5477-
std::move(chunk_reader),
5478-
field_metas,
5479-
use_mmap,
5480-
mmap_config.GetMmapPopulate(),
5481-
mmap_dir_path,
5482-
milvus_field_ids.size(),
5483-
load_info->GetPriority(),
5484-
eager_load,
5485-
warmup_policy,
5486-
cache_key_suffix,
5487-
load_info->GetEstimatedBytesPerRow(),
5488-
load_info->GetInsertChannel());
5473+
std::unique_ptr<Translator<GroupChunk>> translator;
5474+
if (storagev2translator::StorageV2AsyncLoadEnabled()) {
5475+
translator =
5476+
std::make_unique<storagev2translator::ManifestGroupTranslatorV2>(
5477+
get_segment_id(),
5478+
GroupChunkType::DEFAULT,
5479+
index,
5480+
std::move(chunk_reader),
5481+
field_metas,
5482+
use_mmap,
5483+
mmap_config.GetMmapPopulate(),
5484+
mmap_dir_path,
5485+
milvus_field_ids.size(),
5486+
load_info->GetPriority(),
5487+
eager_load,
5488+
warmup_policy,
5489+
cache_key_suffix,
5490+
load_info->GetEstimatedBytesPerRow(),
5491+
load_info->GetInsertChannel());
5492+
} else {
5493+
translator =
5494+
std::make_unique<storagev2translator::ManifestGroupTranslator>(
5495+
get_segment_id(),
5496+
GroupChunkType::DEFAULT,
5497+
index,
5498+
std::move(chunk_reader),
5499+
field_metas,
5500+
use_mmap,
5501+
mmap_config.GetMmapPopulate(),
5502+
mmap_dir_path,
5503+
milvus_field_ids.size(),
5504+
load_info->GetPriority(),
5505+
eager_load,
5506+
warmup_policy,
5507+
cache_key_suffix,
5508+
load_info->GetEstimatedBytesPerRow(),
5509+
load_info->GetInsertChannel());
5510+
}
54895511
auto chunked_column_group =
54905512
std::make_shared<ChunkedColumnGroup>(std::move(translator));
54915513

‎internal/core/src/segcore/storagev2translator/ManifestGroupTranslatorV2Test.cpp‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
#include "pb/common.pb.h"
3333
#include "segcore/storagev2translator/AsyncLoadPipeline.h"
3434
#include "segcore/storagev2translator/ManifestGroupTranslator.h"
35+
#include "segcore/storagev2translator/StorageV2Config.h"
3536
#include "storage/EntryStreamUtils.h"
3637
#include "test_utils/Constants.h"
3738
#include "test_utils/DataGen.h"
@@ -41,6 +42,20 @@ using namespace milvus;
4142
using namespace milvus::segcore;
4243
using namespace milvus::segcore::storagev2translator;
4344

45+
TEST(StorageV2ConfigTest, AsyncLoadDefaultsToEnabledAndCanBeToggled) {
46+
const bool old_enabled = StorageV2AsyncLoadEnabled();
47+
auto restore_enabled = folly::makeGuard(
48+
[old_enabled]() { SetStorageV2AsyncLoadEnabled(old_enabled); });
49+
50+
EXPECT_TRUE(StorageV2AsyncLoadEnabled());
51+
52+
SetStorageV2AsyncLoadEnabled(false);
53+
EXPECT_FALSE(StorageV2AsyncLoadEnabled());
54+
55+
SetStorageV2AsyncLoadEnabled(true);
56+
EXPECT_TRUE(StorageV2AsyncLoadEnabled());
57+
}
58+
4459
class ManifestGroupTranslatorV2ParityTest
4560
: public ::testing::TestWithParam<bool> {
4661
protected:

‎internal/core/src/segcore/storagev2translator/StorageV2Config.h‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
// Licensed to the LF AI & Data foundation under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing, software
12+
// distributed under the License is distributed on an "AS IS" BASIS,
13+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
// See the License for the specific language governing permissions and
15+
// limitations under the License.
16+
17+
#pragma once
18+
19+
#include <atomic>
20+
21+
namespace milvus::segcore::storagev2translator {
22+
23+
// Runtime-configurable via queryNode.segcore.storageV2.enableAsyncLoad.
24+
inline std::atomic<bool>&
25+
storage_v2_async_load_enabled_atomic() {
26+
static std::atomic<bool> instance{true};
27+
return instance;
28+
}
29+
30+
inline bool
31+
StorageV2AsyncLoadEnabled() {
32+
return storage_v2_async_load_enabled_atomic().load(
33+
std::memory_order_acquire);
34+
}
35+
36+
inline void
37+
SetStorageV2AsyncLoadEnabled(bool enabled) {
38+
storage_v2_async_load_enabled_atomic().store(enabled,
39+
std::memory_order_release);
40+
}
41+
42+
} // namespace milvus::segcore::storagev2translator

‎internal/querynodev2/server.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,16 @@ func (node *QueryNode) RegisterSegcoreConfigWatcher() {
286286
mlog.Info(node.ctx, "queryNode.segcore.storageV2.cellTargetSizeBytes updated",
287287
mlog.Int64("bytes", newBytes))
288288
}))
289+
pt.Watch(pt.QueryNodeCfg.StorageV2EnableAsyncLoad.Key,
290+
config.NewHandler("queryNode.segcore.storageV2.enableAsyncLoad", func(evt *config.Event) {
291+
if !evt.HasUpdated {
292+
return
293+
}
294+
enabled := paramtable.Get().QueryNodeCfg.StorageV2EnableAsyncLoad.GetAsBool()
295+
initcore.UpdateStorageV2AsyncLoadEnabled(enabled)
296+
mlog.Info(node.ctx, "queryNode.segcore.storageV2.enableAsyncLoad updated",
297+
mlog.Bool("enabled", enabled))
298+
}))
289299
initcore.RegisterArrowReaderConfigWatchers(pt, "querynode")
290300
}
291301

‎internal/util/initcore/query_node.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -166,6 +166,8 @@ func doInitQueryNodeOnce(ctx context.Context) error {
166166

167167
cStorageV2CellTargetSizeBytes := C.int64_t(paramtable.Get().QueryNodeCfg.StorageV2CellTargetSizeBytes.GetAsInt64())
168168
C.SetStorageV2CellTargetSizeBytes(cStorageV2CellTargetSizeBytes)
169+
cStorageV2EnableAsyncLoad := C.bool(paramtable.Get().QueryNodeCfg.StorageV2EnableAsyncLoad.GetAsBool())
170+
C.SetStorageV2AsyncLoadEnabled(cStorageV2EnableAsyncLoad)
169171
enableParquetStatsSkipIndex := paramtable.Get().CommonCfg.ParquetStatsSkipIndex.GetAsBool()
170172
C.SetDefaultEnableParquetStatsSkipIndex(C.bool(enableParquetStatsSkipIndex))
171173

‎internal/util/initcore/util.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,10 @@ func UpdateStorageV2CellTargetSizeBytes(bytes int64) {
189189
C.SetStorageV2CellTargetSizeBytes(C.int64_t(bytes))
190190
}
191191

192+
func UpdateStorageV2AsyncLoadEnabled(enabled bool) {
193+
C.SetStorageV2AsyncLoadEnabled(C.bool(enabled))
194+
}
195+
192196
func UpdateDefaultGrowingJSONKeyStatsEnable(enable bool) {
193197
C.SetDefaultGrowingJSONKeyStatsEnable(C.bool(enable))
194198
}

‎pkg/util/paramtable/component_param.go‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3671,6 +3671,7 @@ type queryNodeConfig struct {
36713671
// Target average byte size per storage v2 cache cell. Parquet row groups
36723672
// are packed into cells so rgs_per_cell * avg_rg_size ≈ this value.
36733673
StorageV2CellTargetSizeBytes ParamItem `refreshable:"true"`
3674+
StorageV2EnableAsyncLoad ParamItem `refreshable:"true"`
36743675

36753676
EnableWorkerSQCostMetrics ParamItem `refreshable:"true"`
36763677

@@ -4871,6 +4872,15 @@ user-task-polling:
48714872
}
48724873
p.StorageV2CellTargetSizeBytes.Init(base.mgr)
48734874

4875+
p.StorageV2EnableAsyncLoad = ParamItem{
4876+
Key: "queryNode.segcore.storageV2.enableAsyncLoad",
4877+
Version: "3.0.0",
4878+
DefaultValue: "true",
4879+
Doc: "Whether to use the async storage v3 column-group load translator. Disable to compare with the original synchronous translator.",
4880+
Export: true,
4881+
}
4882+
p.StorageV2EnableAsyncLoad.Init(base.mgr)
4883+
48744884
p.EnableWorkerSQCostMetrics = ParamItem{
48754885
Key: "queryNode.enableWorkerSQCostMetrics",
48764886
Version: "2.3.0",

0 commit comments

Comments
 (0)