enhance: Make dataNode.import.maxConcurrentTaskNum dynamic (#37102)

Resize import execution pool when config
`dataNode.import.maxConcurrentTaskNum` update.

issue: https://github.com/milvus-io/milvus/issues/37095

Signed-off-by: bigsheeper <yihao.dai@zilliz.com>
pull/37160/head
yihao.dai 2024-10-25 16:51:29 +08:00 committed by GitHub
parent 1cc9cb49ad
commit d7b2906318
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 93 additions and 2 deletions

View File

@ -17,8 +17,14 @@
package importv2
import (
"context"
"runtime"
"sync"
"go.uber.org/zap"
"github.com/milvus-io/milvus/pkg/config"
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util/conc"
"github.com/milvus-io/milvus/pkg/util/paramtable"
)
@ -29,10 +35,32 @@ var (
)
func initExecPool() {
pt := paramtable.Get()
initPoolSize := paramtable.Get().DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt()
execPool = conc.NewPool[any](
paramtable.Get().DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt(),
conc.WithPreAlloc(true),
initPoolSize,
conc.WithPreAlloc(false), // pre alloc must be false to resize pool dynamically, use warmup to alloc worker here
conc.WithDisablePurge(true),
)
conc.WarmupPool(execPool, runtime.LockOSThread)
watchKey := pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key
pt.Watch(watchKey, config.NewHandler(watchKey, resizeExecPool))
log.Info("init import execution pool done", zap.Int("size", initPoolSize))
}
func resizeExecPool(evt *config.Event) {
if evt.HasUpdated {
newSize := paramtable.Get().DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt()
log := log.Ctx(context.Background()).With(zap.Int("newSize", newSize))
err := GetExecPool().Resize(newSize)
if err != nil {
log.Warn("failed to resize pool", zap.Error(err))
return
}
log.Info("pool resize successfully")
}
}
func GetExecPool() *conc.Pool[any] {

View File

@ -0,0 +1,63 @@
// 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 importv2
import (
"fmt"
"testing"
"github.com/milvus-io/milvus/pkg/config"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/stretchr/testify/assert"
)
func TestResizePools(t *testing.T) {
paramtable.Get().Init(paramtable.NewBaseTable(paramtable.SkipRemote(true)))
pt := paramtable.Get()
defer func() {
_ = pt.Reset(pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key)
}()
t.Run("ExecPool", func(t *testing.T) {
expectedCap := pt.DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt()
assert.Equal(t, expectedCap, GetExecPool().Cap())
resizeExecPool(&config.Event{
HasUpdated: true,
})
assert.Equal(t, expectedCap, GetExecPool().Cap())
_ = pt.Save(pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key, fmt.Sprintf("%d", expectedCap*2))
expectedCap = pt.DataNodeCfg.MaxConcurrentImportTaskNum.GetAsInt()
resizeExecPool(&config.Event{
HasUpdated: true,
})
assert.Equal(t, expectedCap, GetExecPool().Cap())
_ = pt.Save(pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key, "0")
resizeExecPool(&config.Event{
HasUpdated: true,
})
assert.Equal(t, expectedCap, GetExecPool().Cap(), "pool shall not be resized when newSize is 0")
_ = pt.Save(pt.DataNodeCfg.MaxConcurrentImportTaskNum.Key, "invalid")
resizeExecPool(&config.Event{
HasUpdated: true,
})
assert.Equal(t, expectedCap, GetExecPool().Cap())
})
}