From d6982c74c4de5e7f059c79ea83fd1abc59112283 Mon Sep 17 00:00:00 2001 From: "cai.zhang" Date: Fri, 17 Jul 2026 17:38:40 +0800 Subject: [PATCH] enhance: record and expose DataNode task cost_time / cost_cpu_num via QueryTask (#51443) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit issue: #49156 ## Summary - Record `cost_time` / `cost_cpu_num` on DataNode index build tasks and surface them through `QueryTask` response `Properties`. - Reserve the unified `cost_time` / `cost_cpu_num` keys under `pkg/taskcommon` so future PRs can wire up the other `CreateTask`-dispatched task types. - Introduce small `internal/datanode/taskcost` helper (`NowMs`, `ElapsedMs`, `EstimateIndexBuildCPUNum`) for index build task cost accounting. - Final state/failReason and execution-end cost are written in a single `TaskManager` critical section, and `QueryTask` reads them from a single snapshot, so state and cost in one response are always consistent. ## Notes This PR is producer-side only. DataCoord / proxy / metrics exporter do not consume these properties yet; downstream consumption should land in a follow-up linked from #49156. `cost_cpu_num` is a coarse, statistics-only estimate for observability (vector index ≈ knowhere build pool size = NumCPU in cluster mode; scalar index = 1). It is not used for coordinator-side scheduling, so the standalone pool-ratio deviation (NumCPU × buildIndexThreadPoolRatio) and pool sharing across concurrent builds are intentionally not modeled. Supersedes #49157 (that PR's head branch was rebased onto latest master and force-pushed, so GitHub no longer allows reopening it). Content is identical, rebased to a single commit on top of current master. ## Test plan - [x] `go test ./taskcommon/...` from `pkg/` - [x] `go test -tags dynamic,test -gcflags="all=-N -l" -count=1 ./internal/datanode/taskcost` - [ ] `go test -tags dynamic,test -gcflags="all=-N -l" -count=1 ./internal/datanode/index -run 'TestIndexTaskSchedulerRecordsIndexTaskCost|Test_statsTaskInfoSuite/Test_IndexTaskCostMethods'` - [ ] Manual sanity: `QueryTask` returns `cost_time` / `cost_cpu_num` for DataNode index build task success and failure paths. --------- Signed-off-by: cai.zhang Co-authored-by: Claude Opus 4.7 --- internal/datanode/index/scheduler.go | 45 ++++++++- internal/datanode/index/scheduler_test.go | 100 +++++++++++++++++++ internal/datanode/index/task_index.go | 13 +++ internal/datanode/index/task_manager.go | 49 +++++++++ internal/datanode/index/task_manager_test.go | 29 ++++++ internal/datanode/services.go | 32 ++++-- internal/datanode/services_test.go | 24 +++++ internal/datanode/taskcost/cost.go | 60 +++++++++++ internal/datanode/taskcost/cost_test.go | 49 +++++++++ pkg/taskcommon/properties.go | 36 ++++++- pkg/taskcommon/properties_test.go | 26 +++++ 11 files changed, 448 insertions(+), 15 deletions(-) create mode 100644 internal/datanode/taskcost/cost.go create mode 100644 internal/datanode/taskcost/cost_test.go diff --git a/internal/datanode/index/scheduler.go b/internal/datanode/index/scheduler.go index 8dddfbe6d5..cd63b42bf0 100644 --- a/internal/datanode/index/scheduler.go +++ b/internal/datanode/index/scheduler.go @@ -21,10 +21,12 @@ import ( "context" "runtime/debug" "sync" + "time" "github.com/cockroachdb/errors" "go.uber.org/atomic" + "github.com/milvus-io/milvus/internal/datanode/taskcost" "github.com/milvus-io/milvus/internal/storagev2" "github.com/milvus-io/milvus/pkg/v3/mlog" "github.com/milvus-io/milvus/pkg/v3/proto/indexpb" @@ -242,19 +244,52 @@ func (sched *TaskScheduler) processTask(t Task) { }() sched.TaskQueue.AddActiveTask(t) defer sched.TaskQueue.PopActiveTask(t.Name()) - mlog.Debug(t.Ctx(), "process task", mlog.String("task", t.Name())) + var ( + indexTask *indexBuildTask + costCPUNum int64 + execStart time.Time + ) + if ibt, ok := t.(*indexBuildTask); ok { + indexTask = ibt + costCPUNum = taskcost.EstimateIndexBuildCPUNum(indexTask.IsVectorIndex()) + // execStart carries a monotonic clock reading; CostTimeMs derived from + // it is immune to wall-clock steps. ExecStartMs/ExecEndMs stay wall-clock + // timestamps for external exposure. + execStart = time.Now() + indexTask.manager.StoreIndexTaskExecutionStart(indexTask.req.GetClusterID(), indexTask.req.GetBuildID(), taskcost.NowMs(), costCPUNum) + mlog.Debug(t.Ctx(), "process task", mlog.String("task", t.Name()), mlog.Int64("costCPUNum", costCPUNum)) + } else { + mlog.Debug(t.Ctx(), "process task", mlog.String("task", t.Name())) + } + pipelines := []func(context.Context) error{t.PreExecute, t.Execute, t.PostExecute} for _, fn := range pipelines { if err := wrap(fn); err != nil { - mlog.Warn(t.Ctx(), "process task failed", mlog.Err(err)) - t.SetState(getStateFromError(err), err.Error()) + if indexTask != nil { + costTimeMs := taskcost.ElapsedMs(execStart) + // End bookkeeping and final state must land in one critical + // section, so a concurrent QueryTask never sees a final cost + // paired with an in-progress state. + indexTask.SetStateWithCost(getStateFromError(err), err.Error(), taskcost.NowMs(), costTimeMs) + mlog.Warn(t.Ctx(), "process task failed", mlog.Err(err), mlog.Int64("costTimeMs", costTimeMs), mlog.Int64("costCPUNum", costCPUNum)) + } else { + t.SetState(getStateFromError(err), err.Error()) + mlog.Warn(t.Ctx(), "process task failed", mlog.Err(err)) + } return } } - t.SetState(indexpb.JobState_JobStateFinished, "") + if indexTask != nil { + costTimeMs := taskcost.ElapsedMs(execStart) + indexTask.SetStateWithCost(indexpb.JobState_JobStateFinished, "", taskcost.NowMs(), costTimeMs) + mlog.Debug(t.Ctx(), "process task completed", mlog.String("task", t.Name()), mlog.Int64("costTimeMs", costTimeMs), mlog.Int64("costCPUNum", costCPUNum)) + } else { + t.SetState(indexpb.JobState_JobStateFinished, "") + mlog.Debug(t.Ctx(), "process task completed", mlog.String("task", t.Name())) + } // Publish filesystem metrics after index task completion - if indexTask, ok := t.(*indexBuildTask); ok { + if indexTask != nil { if indexTask.req != nil && indexTask.req.GetStorageConfig() != nil { storagev2.PublishFilesystemMetricsWithConfig(indexTask.req.GetStorageConfig()) } diff --git a/internal/datanode/index/scheduler_test.go b/internal/datanode/index/scheduler_test.go index dd9be81148..4d924f2de4 100644 --- a/internal/datanode/index/scheduler_test.go +++ b/internal/datanode/index/scheduler_test.go @@ -23,10 +23,16 @@ import ( "testing" "time" + "github.com/bytedance/mockey" "github.com/cockroachdb/errors" "github.com/stretchr/testify/assert" + "github.com/milvus-io/milvus-proto/go-api/v3/commonpb" + "github.com/milvus-io/milvus-proto/go-api/v3/schemapb" + "github.com/milvus-io/milvus/pkg/v3/common" "github.com/milvus-io/milvus/pkg/v3/proto/indexpb" + "github.com/milvus-io/milvus/pkg/v3/proto/workerpb" + "github.com/milvus-io/milvus/pkg/v3/util/hardware" "github.com/milvus-io/milvus/pkg/v3/util/paramtable" ) @@ -230,3 +236,97 @@ func TestIndexTaskScheduler(t *testing.T) { assert.Equal(t, task.GetState(), indexpb.JobState_JobStateFinished) } } + +func newSchedulerIndexBuildTask(t *testing.T, manager *TaskManager, buildID int64) *indexBuildTask { + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) + + req := &workerpb.CreateJobRequest{ + ClusterID: "test-cluster", + BuildID: buildID, + IndexParams: []*commonpb.KeyValuePair{ + {Key: common.IndexTypeKey, Value: "STL_SORT"}, + }, + Field: &schemapb.FieldSchema{ + FieldID: 100, + DataType: schemapb.DataType_Int64, + }, + } + manager.LoadOrStoreIndexTask(req.GetClusterID(), req.GetBuildID(), &IndexTaskInfo{ + State: commonpb.IndexState_InProgress, + }) + return NewIndexBuildTask(ctx, cancel, req, nil, manager, nil) +} + +func TestIndexTaskSchedulerRecordsIndexTaskCost(t *testing.T) { + paramtable.Init() + + t.Run("success records execution cost", func(t *testing.T) { + manager := NewTaskManager(context.Background()) + task := newSchedulerIndexBuildTask(t, manager, 1001) + + preMock := mockey.Mock((*indexBuildTask).PreExecute).Return(nil).Build() + defer preMock.UnPatch() + executeMock := mockey.Mock((*indexBuildTask).Execute).Return(nil).Build() + defer executeMock.UnPatch() + postMock := mockey.Mock((*indexBuildTask).PostExecute).Return(nil).Build() + defer postMock.UnPatch() + + scheduler := NewTaskScheduler(context.Background()) + scheduler.processTask(task) + + info := manager.GetIndexTaskInfo("test-cluster", 1001) + assert.NotNil(t, info) + assert.Equal(t, commonpb.IndexState_Finished, info.State) + assert.Greater(t, info.ExecStartMs, int64(0)) + assert.GreaterOrEqual(t, info.ExecEndMs, info.ExecStartMs) + assert.GreaterOrEqual(t, info.CostTimeMs, int64(0)) + assert.Equal(t, int64(1), info.CostCPUNum) + }) + + t.Run("pre execute failure still records execution end", func(t *testing.T) { + manager := NewTaskManager(context.Background()) + task := newSchedulerIndexBuildTask(t, manager, 1002) + expectedErr := errors.New("pre execute failed") + + preMock := mockey.Mock((*indexBuildTask).PreExecute).Return(expectedErr).Build() + defer preMock.UnPatch() + + scheduler := NewTaskScheduler(context.Background()) + scheduler.processTask(task) + + info := manager.GetIndexTaskInfo("test-cluster", 1002) + assert.NotNil(t, info) + assert.Equal(t, commonpb.IndexState_Retry, info.State) + assert.Equal(t, expectedErr.Error(), info.FailReason) + assert.Greater(t, info.ExecStartMs, int64(0)) + assert.GreaterOrEqual(t, info.ExecEndMs, info.ExecStartMs) + assert.GreaterOrEqual(t, info.CostTimeMs, int64(0)) + assert.Equal(t, int64(1), info.CostCPUNum) + }) + + t.Run("vector index records build pool cpu num", func(t *testing.T) { + manager := NewTaskManager(context.Background()) + task := newSchedulerIndexBuildTask(t, manager, 1003) + + vecMock := mockey.Mock((*indexBuildTask).IsVectorIndex).Return(true).Build() + defer vecMock.UnPatch() + preMock := mockey.Mock((*indexBuildTask).PreExecute).Return(nil).Build() + defer preMock.UnPatch() + executeMock := mockey.Mock((*indexBuildTask).Execute).Return(nil).Build() + defer executeMock.UnPatch() + postMock := mockey.Mock((*indexBuildTask).PostExecute).Return(nil).Build() + defer postMock.UnPatch() + + scheduler := NewTaskScheduler(context.Background()) + scheduler.processTask(task) + + info := manager.GetIndexTaskInfo("test-cluster", 1003) + assert.NotNil(t, info) + assert.Equal(t, commonpb.IndexState_Finished, info.State) + assert.Greater(t, info.ExecStartMs, int64(0)) + assert.GreaterOrEqual(t, info.ExecEndMs, info.ExecStartMs) + assert.GreaterOrEqual(t, info.CostTimeMs, int64(0)) + assert.Equal(t, int64(hardware.GetCPUNum()), info.CostCPUNum) + }) +} diff --git a/internal/datanode/index/task_index.go b/internal/datanode/index/task_index.go index b7af6edfb2..ee8bef1f04 100644 --- a/internal/datanode/index/task_index.go +++ b/internal/datanode/index/task_index.go @@ -118,6 +118,19 @@ func (it *indexBuildTask) Name() string { func (it *indexBuildTask) SetState(state indexpb.JobState, failReason string) { it.manager.StoreIndexTaskState(it.req.GetClusterID(), it.req.GetBuildID(), commonpb.IndexState(state), failReason) + it.observeStateMetrics(state) +} + +// SetStateWithCost stores the final state/failReason together with the +// execution-end cost bookkeeping in one TaskManager critical section, so a +// concurrent QueryTask can never observe a final cost paired with a stale +// in-progress state. +func (it *indexBuildTask) SetStateWithCost(state indexpb.JobState, failReason string, endMs, costTimeMs int64) { + it.manager.StoreIndexTaskExecutionEndWithState(it.req.GetClusterID(), it.req.GetBuildID(), endMs, costTimeMs, commonpb.IndexState(state), failReason) + it.observeStateMetrics(state) +} + +func (it *indexBuildTask) observeStateMetrics(state indexpb.JobState) { if state == indexpb.JobState_JobStateFinished { metrics.DataNodeBuildIndexLatency.WithLabelValues(paramtable.GetStringNodeID()).Observe(it.tr.ElapseSpan().Seconds()) metrics.DataNodeIndexTaskLatencyInQueue.WithLabelValues(paramtable.GetStringNodeID()).Observe(float64(it.queueDur.Milliseconds())) diff --git a/internal/datanode/index/task_manager.go b/internal/datanode/index/task_manager.go index c2aa47acc5..0718697eeb 100644 --- a/internal/datanode/index/task_manager.go +++ b/internal/datanode/index/task_manager.go @@ -44,6 +44,10 @@ type IndexTaskInfo struct { CurrentIndexVersion int32 CurrentScalarIndexVersion int32 IndexStorePathVersion indexpb.IndexStorePathVersion + ExecStartMs int64 + ExecEndMs int64 + CostTimeMs int64 + CostCPUNum int64 // task statistics statistic *indexpb.JobInfo @@ -60,6 +64,10 @@ func (i *IndexTaskInfo) Clone() *IndexTaskInfo { CurrentIndexVersion: i.CurrentIndexVersion, CurrentScalarIndexVersion: i.CurrentScalarIndexVersion, IndexStorePathVersion: i.IndexStorePathVersion, + ExecStartMs: i.ExecStartMs, + ExecEndMs: i.ExecEndMs, + CostTimeMs: i.CostTimeMs, + CostCPUNum: i.CostCPUNum, statistic: typeutil.Clone(i.statistic), } } @@ -130,6 +138,37 @@ func (m *TaskManager) StoreIndexTaskState(ClusterID string, buildID typeutil.Uni } } +func (m *TaskManager) StoreIndexTaskExecutionStart(clusterID string, buildID typeutil.UniqueID, startMs int64, costCPUNum int64) { + key := Key{ClusterID: clusterID, TaskID: buildID} + m.stateLock.Lock() + defer m.stateLock.Unlock() + if task, ok := m.indexTasks[key]; ok { + task.ExecStartMs = startMs + task.CostCPUNum = costCPUNum + task.ExecEndMs = 0 + task.CostTimeMs = 0 + } +} + +// StoreIndexTaskExecutionEndWithState records the execution-end cost bookkeeping +// together with the final task state/failReason in a single critical section, so +// a concurrent reader can never observe a final cost paired with a stale state +// (or vice versa). +func (m *TaskManager) StoreIndexTaskExecutionEndWithState(clusterID string, buildID typeutil.UniqueID, endMs int64, costTimeMs int64, state commonpb.IndexState, failReason string) { + key := Key{ClusterID: clusterID, TaskID: buildID} + m.stateLock.Lock() + defer m.stateLock.Unlock() + if task, ok := m.indexTasks[key]; ok { + mlog.Debug(m.ctx, "store task execution end with state", mlog.String("clusterID", clusterID), mlog.FieldBuildID(buildID), + mlog.String("state", state.String()), mlog.String("fail reason", failReason), + mlog.Int64("costTimeMs", costTimeMs)) + task.ExecEndMs = endMs + task.CostTimeMs = costTimeMs + task.State = state + task.FailReason = failReason + } +} + func (m *TaskManager) ForeachIndexTaskInfo(fn func(ClusterID string, buildID typeutil.UniqueID, info *IndexTaskInfo)) { m.stateLock.Lock() defer m.stateLock.Unlock() @@ -138,6 +177,16 @@ func (m *TaskManager) ForeachIndexTaskInfo(fn func(ClusterID string, buildID typ } } +func (m *TaskManager) GetIndexTaskInfo(clusterID string, buildID typeutil.UniqueID) *IndexTaskInfo { + m.stateLock.Lock() + defer m.stateLock.Unlock() + + if info, ok := m.indexTasks[Key{ClusterID: clusterID, TaskID: buildID}]; ok { + return info.Clone() + } + return nil +} + func (m *TaskManager) StoreIndexFilesAndStatistic( ClusterID string, buildID typeutil.UniqueID, diff --git a/internal/datanode/index/task_manager_test.go b/internal/datanode/index/task_manager_test.go index 366f326bb7..79dd237fd7 100644 --- a/internal/datanode/index/task_manager_test.go +++ b/internal/datanode/index/task_manager_test.go @@ -156,3 +156,32 @@ func (s *statsTaskInfoSuite) TestIndexTaskInfoReturnsIndexStorePathVersion() { cloned := info.Clone() s.Equal(indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_COLLECTION_ROOTED, cloned.IndexStorePathVersion) } + +func (s *statsTaskInfoSuite) Test_IndexTaskCostMethods() { + buildID := int64(200) + s.manager.LoadOrStoreIndexTask(s.cluster, buildID, &IndexTaskInfo{State: commonpb.IndexState_InProgress}) + + s.manager.StoreIndexTaskExecutionStart(s.cluster, buildID, 111, 4) + s.manager.StoreIndexTaskExecutionEndWithState(s.cluster, buildID, 222, 111, commonpb.IndexState_Failed, "mock failure") + + info := s.manager.GetIndexTaskInfo(s.cluster, buildID) + s.Require().NotNil(info) + s.Equal(int64(111), info.ExecStartMs) + s.Equal(int64(222), info.ExecEndMs) + s.Equal(int64(111), info.CostTimeMs) + s.Equal(int64(4), info.CostCPUNum) + // state/failReason land together with the cost in one critical section + s.Equal(commonpb.IndexState_Failed, info.State) + s.Equal("mock failure", info.FailReason) + + info.CostCPUNum = 999 + reloaded := s.manager.GetIndexTaskInfo(s.cluster, buildID) + s.Equal(int64(4), reloaded.CostCPUNum) + + s.manager.StoreIndexTaskExecutionStart(s.cluster, buildID, 333, 8) + reloaded = s.manager.GetIndexTaskInfo(s.cluster, buildID) + s.Equal(int64(333), reloaded.ExecStartMs) + s.Equal(int64(8), reloaded.CostCPUNum) + s.Equal(int64(0), reloaded.ExecEndMs) + s.Equal(int64(0), reloaded.CostTimeMs) +} diff --git a/internal/datanode/services.go b/internal/datanode/services.go index 225e3542fe..7309996243 100644 --- a/internal/datanode/services.go +++ b/internal/datanode/services.go @@ -814,15 +814,31 @@ func (node *DataNode) QueryTask(ctx context.Context, request *workerpb.QueryTask } return wrapQueryTaskResult(resp, resProperties) case taskcommon.Index: - resp, err := node.queryIndexTask(ctx, &workerpb.QueryJobsRequest{ClusterID: clusterID, TaskIDs: []int64{taskID}}) - if err != nil { - return nil, err - } + // State/reason and cost must come from one snapshot cloned under one + // lock, so a concurrently completing task can never yield a final cost + // paired with an in-progress state (or vice versa). resProperties := taskcommon.NewProperties(nil) - results := resp.GetIndexJobResults().GetResults() - if len(results) > 0 { - resProperties.AppendTaskState(taskcommon.State(results[0].GetState())) - resProperties.AppendReason(results[0].GetFailReason()) + info := node.taskManager.GetIndexTaskInfo(clusterID, taskID) + if info == nil { + resProperties.AppendCostTime(0) + resProperties.AppendCostCPUNum(0) + resp := &workerpb.QueryJobsV2Response{ + Status: merr.Status(merr.WrapErrServiceInternalMsg("tasks '%v' not found", []int64{taskID})), + } + return wrapQueryTaskResult(resp, resProperties) + } + resProperties.AppendTaskState(taskcommon.State(info.State)) + resProperties.AppendReason(info.FailReason) + resProperties.AppendCostTime(info.CostTimeMs) + resProperties.AppendCostCPUNum(info.CostCPUNum) + resp := &workerpb.QueryJobsV2Response{ + Status: merr.Success(), + ClusterID: clusterID, + Result: &workerpb.QueryJobsV2Response_IndexJobResults{ + IndexJobResults: &workerpb.IndexJobResults{ + Results: []*workerpb.IndexTaskInfo{info.ToIndexTaskInfo(taskID)}, + }, + }, } return wrapQueryTaskResult(resp, resProperties) case taskcommon.Stats: diff --git a/internal/datanode/services_test.go b/internal/datanode/services_test.go index 11bd5ddcf9..b706f28697 100644 --- a/internal/datanode/services_test.go +++ b/internal/datanode/services_test.go @@ -34,6 +34,7 @@ import ( "github.com/milvus-io/milvus/internal/compaction" "github.com/milvus-io/milvus/internal/datanode/compactor" "github.com/milvus-io/milvus/internal/datanode/external" + "github.com/milvus-io/milvus/internal/datanode/index" "github.com/milvus-io/milvus/internal/storage" "github.com/milvus-io/milvus/internal/util/sessionutil" "github.com/milvus-io/milvus/pkg/v3/common" @@ -678,6 +679,29 @@ func (s *DataNodeServicesSuite) TestQueryTask() { s.True(strings.Contains(resp.GetStatus().GetReason(), "not found")) }) + s.Run("query index task with cost", func() { + s.node.taskManager.LoadOrStoreIndexTask("cluster-0", 101, &index.IndexTaskInfo{State: commonpb.IndexState_InProgress}) + s.node.taskManager.StoreIndexTaskExecutionStart("cluster-0", 101, 100, 3) + s.node.taskManager.StoreIndexTaskExecutionEndWithState("cluster-0", 101, 180, 80, commonpb.IndexState_Finished, "") + + req := &workerpb.QueryTaskRequest{ + Properties: map[string]string{ + taskcommon.ClusterIDKey: "cluster-0", + taskcommon.TypeKey: taskcommon.Index, + taskcommon.TaskIDKey: "101", + }, + } + resp, err := s.node.QueryTask(s.ctx, req) + s.NoError(merr.CheckRPCCall(resp, err)) + props := taskcommon.NewProperties(resp.GetProperties()) + // state and cost come from the same snapshot + state, err := props.GetTaskState() + s.NoError(err) + s.Equal(taskcommon.State(commonpb.IndexState_Finished), state) + s.Equal(int64(80), props.GetCostTime()) + s.Equal(int64(3), props.GetCostCPUNum()) + }) + s.Run("invalid task type", func() { req := &workerpb.QueryTaskRequest{ Properties: map[string]string{ diff --git a/internal/datanode/taskcost/cost.go b/internal/datanode/taskcost/cost.go new file mode 100644 index 0000000000..9e505c8251 --- /dev/null +++ b/internal/datanode/taskcost/cost.go @@ -0,0 +1,60 @@ +// 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 taskcost + +import ( + "time" + + "github.com/milvus-io/milvus/pkg/v3/util/hardware" +) + +// NowMs returns the current wall-clock time in milliseconds. Use it only for +// externally exposed timestamps (ExecStartMs / ExecEndMs); never subtract two +// NowMs readings to derive a duration — use ElapsedMs instead. +func NowMs() int64 { + return time.Now().UnixMilli() +} + +// ElapsedMs returns the duration since start in milliseconds. It relies on the +// monotonic clock reading embedded in start (time.Since), so the result is +// immune to wall-clock steps such as NTP corrections and is never negative. +func ElapsedMs(start time.Time) int64 { + return time.Since(start).Milliseconds() +} + +// EstimateIndexBuildCPUNum returns the approximate number of CPU threads an +// index build task consumes during execution. +// +// The value is task-level statistics for observability only and is never used +// for coordinator-side scheduling, so approximation error is acceptable (e.g. +// standalone sizes the build pool to NumCPU * buildIndexThreadPoolRatio, and +// concurrent builds share one pool). +// +// Vector indexes go through knowhere's build thread pool (sized to NumCPU +// in cluster mode, see internal/datanode/index/init_segcore.go:74), so they +// report hardware.GetCPUNum(). All current scalar indexes build +// single-threaded: the tantivy wrapper hardcodes 1 thread +// (internal/core/thirdparty/tantivy/tantivy-wrapper.h:23 — covers INVERTED / +// NGRAM / TEXT_MATCH / JSON_INVERTED), and the rest (BITMAP / STL_SORT / +// STRING_SORT / MARISA / HYBRID / RTREE) are plain loops with no thread +// pool / OpenMP / std::async. They therefore report 1. +func EstimateIndexBuildCPUNum(isVectorIndex bool) int64 { + if isVectorIndex { + return int64(hardware.GetCPUNum()) + } + return 1 +} diff --git a/internal/datanode/taskcost/cost_test.go b/internal/datanode/taskcost/cost_test.go new file mode 100644 index 0000000000..b11626df76 --- /dev/null +++ b/internal/datanode/taskcost/cost_test.go @@ -0,0 +1,49 @@ +// 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 taskcost + +import ( + "testing" + "time" + + "github.com/stretchr/testify/assert" + + "github.com/milvus-io/milvus/pkg/v3/util/hardware" +) + +func TestNowMs(t *testing.T) { + before := time.Now().UnixMilli() + got := NowMs() + after := time.Now().UnixMilli() + assert.GreaterOrEqual(t, got, before) + assert.LessOrEqual(t, got, after) +} + +func TestElapsedMs(t *testing.T) { + start := time.Now() + time.Sleep(10 * time.Millisecond) + got := ElapsedMs(start) + // Sleep guarantees at least 10ms elapsed on the monotonic clock. + assert.GreaterOrEqual(t, got, int64(10)) +} + +func TestEstimateIndexBuildCPUNum(t *testing.T) { + assert.Equal(t, int64(hardware.GetCPUNum()), EstimateIndexBuildCPUNum(true), + "vector index saturates knowhere build pool") + assert.Equal(t, int64(1), EstimateIndexBuildCPUNum(false), + "all scalar indexes currently build single-threaded") +} diff --git a/pkg/taskcommon/properties.go b/pkg/taskcommon/properties.go index f87ab051ee..ae519a617c 100644 --- a/pkg/taskcommon/properties.go +++ b/pkg/taskcommon/properties.go @@ -37,8 +37,10 @@ const ( CollectionIDKey = "collection_id" // result - StateKey = "task_state" - ReasonKey = "task_reason" + StateKey = "task_state" + ReasonKey = "task_reason" + CostTimeKey = "cost_time" + CostCPUNumKey = "cost_cpu_num" ) type Properties map[string]string @@ -102,6 +104,14 @@ func (p Properties) AppendTaskState(state State) { p[StateKey] = state.String() } +func (p Properties) AppendCostTime(costTime int64) { + p[CostTimeKey] = fmt.Sprintf("%d", costTime) +} + +func (p Properties) AppendCostCPUNum(costCPUNum int64) { + p[CostCPUNumKey] = fmt.Sprintf("%d", costCPUNum) +} + func (p Properties) GetTaskType() (Type, error) { if _, ok := p[TypeKey]; !ok { return "", WrapErrTaskPropertyLack(TypeKey, p[TaskIDKey]) @@ -206,3 +216,25 @@ func (p Properties) GetCollectionID() (int64, error) { } return collectionID, nil } + +func (p Properties) GetCostTime() int64 { + if _, ok := p[CostTimeKey]; !ok { + return 0 + } + costTime, err := strconv.ParseInt(p[CostTimeKey], 10, 64) + if err != nil { + return 0 + } + return costTime +} + +func (p Properties) GetCostCPUNum() int64 { + if _, ok := p[CostCPUNumKey]; !ok { + return 0 + } + costCPUNum, err := strconv.ParseInt(p[CostCPUNumKey], 10, 64) + if err != nil { + return 0 + } + return costCPUNum +} diff --git a/pkg/taskcommon/properties_test.go b/pkg/taskcommon/properties_test.go index 2d3a9cb487..a899bdbc37 100644 --- a/pkg/taskcommon/properties_test.go +++ b/pkg/taskcommon/properties_test.go @@ -100,3 +100,29 @@ func TestProperties_CollectionID(t *testing.T) { assert.Equal(t, int64(-1), collectionID) }) } + +func TestProperties_CostFields(t *testing.T) { + t.Run("append and get", func(t *testing.T) { + props := NewProperties(nil) + props.AppendCostTime(123) + props.AppendCostCPUNum(4) + + assert.Equal(t, int64(123), props.GetCostTime()) + assert.Equal(t, int64(4), props.GetCostCPUNum()) + }) + + t.Run("missing keys", func(t *testing.T) { + props := NewProperties(nil) + assert.Equal(t, int64(0), props.GetCostTime()) + assert.Equal(t, int64(0), props.GetCostCPUNum()) + }) + + t.Run("invalid values", func(t *testing.T) { + props := NewProperties(map[string]string{ + CostTimeKey: "abc", + CostCPUNumKey: "xyz", + }) + assert.Equal(t, int64(0), props.GetCostTime()) + assert.Equal(t, int64(0), props.GetCostCPUNum()) + }) +}