Skip to content

Commit f986d60

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

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
@@ -581,6 +581,7 @@ queryNode:
581581
deletePoolSizeFactor: 1 # size factor (CPUNum * factor) of the DeleteBatch dispatch pool
582582
storageV2:
583583
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.
584+
enableAsyncLoad: true # Enable the storage v3 async column-group load translator. Disable to compare with the original synchronous translator.
584585
deleteDumpBatchSize: 10000 # Batch size for delete snapshot dump in segcore.
585586
loadMemoryUsageFactor: 1 # The multiply factor of calculating the memory usage while loading segments
586587
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"
@@ -5471,23 +5472,44 @@ ChunkedSegmentSealedImpl::LoadColumnGroup(
54715472
cache_key_suffix = std::to_string(milvus_field_ids.front().get());
54725473
}
54735474

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

‎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
@@ -3672,6 +3672,7 @@ type queryNodeConfig struct {
36723672
// Target average byte size per storage v2 cache cell. Parquet row groups
36733673
// are packed into cells so rgs_per_cell * avg_rg_size ≈ this value.
36743674
StorageV2CellTargetSizeBytes ParamItem `refreshable:"true"`
3675+
StorageV2EnableAsyncLoad ParamItem `refreshable:"true"`
36753676

36763677
EnableWorkerSQCostMetrics ParamItem `refreshable:"true"`
36773678

@@ -4899,6 +4900,15 @@ user-task-polling:
48994900
}
49004901
p.StorageV2CellTargetSizeBytes.Init(base.mgr)
49014902

4903+
p.StorageV2EnableAsyncLoad = ParamItem{
4904+
Key: "queryNode.segcore.storageV2.enableAsyncLoad",
4905+
Version: "3.0.0",
4906+
DefaultValue: "true",
4907+
Doc: "Whether to use the async storage v3 column-group load translator. Disable to compare with the original synchronous translator.",
4908+
Export: true,
4909+
}
4910+
p.StorageV2EnableAsyncLoad.Init(base.mgr)
4911+
49024912
p.EnableWorkerSQCostMetrics = ParamItem{
49034913
Key: "queryNode.enableWorkerSQCostMetrics",
49044914
Version: "2.3.0",

0 commit comments

Comments
 (0)