milvus-io/milvus
1
0
Fork 0
mirror of https://github.com/milvus-io/milvus.git synced 2026-10-11 22:49:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

enhance: add QueryCoord autoscale load resource precheck (#51218)

Browse Source 
## Summary

issue: #51217

Load requests introduce resource demand before QueryNode waterlines
reflect the new data. This PR adds optional QueryCoord resource
prechecks before LoadCollection/LoadPartitions configuration broadcasts
and publishes cumulative demand counters for external autoscaling
decisions.

Design: [Autoscale-Aware Load Resource Admission in
QueryCoord](https://github.com/sijie-ni-0214/milvus/blob/querycoord-autoscale-kernel/docs/design-docs/design_docs/20260721-querycoord-autoscale-load-resource-precheck.md).

## Admission behavior

- Estimate incremental final sealed-segment memory and local-disk demand
per resource group from the requested and current load configurations.
- With tiered eviction disabled, compare demand with memory availability
and each Worker's configured `LOCAL_STORAGE_SIZE` quota minus cached
local-directory usage after admission thresholds. Optionally admit
shortages within configured global autoscale upper bounds.
- With tiered eviction enabled, retain demand estimation and skip QN
capacity sampling, node availability checks, and global autoscale-limit
checks. Existing evictable caches make physical usage unsuitable for
this admission decision; Worker QN retains its loading-stage guard.
- Fail open on metadata, estimation, and capacity-metrics failures.
Capacity-metrics failures preserve successfully estimated demand;
metadata or estimation failures leave no demand to publish.
- Return a resource-insufficient status with `suggested_expand_percent`
when a complete capacity evaluation rejects the request. Proxy preserves
the status and ExtraInfo.
- Admission remains best-effort, without reservations or accounting for
concurrent accepted requests. QueryCoord does not subtract Worker
in-flight load reservations; the Worker guard remains authoritative. The
existing replica node-count prerequisite remains; node provisioning and
scale-from-zero are outside this PR.

## Load demand counters

Expose cumulative counters labeled by resource group:

- `milvus_querycoord_load_demand_memory_bytes`
- `milvus_querycoord_load_demand_disk_bytes`

Counters increase only when demand estimation succeeds, the request
proceeds, and its configuration broadcast succeeds. This includes tiered
loads and loads admitted through capacity-metrics fail-open. Disabled
prechecks, unchanged configurations, zero incremental demand, estimation
failures, explicit rejections, and failed broadcasts do not increase
counters. Load completion and release do not decrement them.

Series start at zero for existing resource groups and are removed on RG
deletion. The dashboard displays cumulative values in MiB; the control
plane computes interval increments, handles counter resets, and combines
demand with QN waterlines and provisioning state. These counters are
observations, not durable reservations. Internal handoff and balancing
tasks do not contribute to them.

## Resource estimation and sampling

- Share segment/index estimation in
`internal/util/segcore/loadresource`, retaining separate final-resource
and loading-peak APIs. Worker QN uses its CGo pool; QueryCoord estimates
sequentially within each request.
- Fetch recovery metadata, full segment metadata, and batched index
metadata. Support Storage V3 resource summaries and dropped-field
metadata. Filter indexes by loaded fields, retaining multiple JSON-path
indexes on one field.
- Apply the load-request collection mmap inheritance logic to cloned
expected/current schemas while preserving field overrides. The final
estimator applies cache ratios only with tiered eviction enabled.
Delegator-side Bloom/BM25 resources are outside the Worker estimate.
- Deduplicate QN resource sampling by node ID, use at most 16 concurrent
requests and one attempt per node, and bound the entire sampling round
to ten seconds. Reuse the snapshot for RG and global calculations.
Exclude embedded SQNs from global capacity and sampling; empty
eligible-node sets are unevaluable and fail open.
- Worker QNs report `LOCAL_STORAGE_SIZE` and cached local-directory
usage through `system_info`. QueryCoord uses these byte values for disk
availability, RG suggestions, and the global disk-capacity sum;
filesystem-wide disk metrics remain for monitoring. Shared filesystems
require per-QN quotas consistent with the shared capacity.

## Configuration

- `queryCoord.autoscale.precheckEnabled`: enable demand estimation and
precheck; default `false`.
- `queryCoord.autoscale.enabled`: allow admission within global
autoscale bounds when current capacity is insufficient; default `false`,
effective only with precheck enabled and tiered eviction disabled.
- `queryCoord.autoscale.maxMemoryLimit`: maximum global Worker QN
physical memory capacity after scaling, in GiB; default `0` means no
additional memory autoscale capacity.
- `queryCoord.autoscale.maxDiskLimit`: maximum sum of Worker QN
`LOCAL_STORAGE_SIZE` quotas after scaling, in GiB; default `0` means no
additional disk autoscale capacity. Both limits use `1 GiB = 1024^3
bytes`.

For non-tiered deployments, configure QueryCoord's memory high-water
level no higher than the Worker admission threshold. This relationship
is a rollout prerequisite and is not enforced by this PR.

## Validation

Local autoscale estimator and targeted QueryCoord tests passed,
including multiple-index filtering, collection mmap overrides, the
ten-second sampling timeout, empty-node fail-open, embedded-SQN
exclusion, tiered demand preservation, capacity-metrics failure
behavior, and local-storage quota accounting for QueryNodes sharing a
filesystem.

Both load callbacks are tested for precheck rejection, broadcast
failure, successful counter updates, tiered loads, and capacity-metrics
failures followed by successful or failed broadcasts. Tests assert
counters remain unchanged during Broadcast and increment only after it
succeeds. These are local unit/callback checks, not end-to-end external
autoscaler validation.

Local `make static-check` passed across core, pkg, client, and go_client
with zero issues, and segcore error-boundary checks passed. Targeted
QueryCoord and QueryNode tests, YAML consistency, and `make build-go`
passed. External autoscaler behavior remains subject to CI and rollout
validation.

---------

Signed-off-by: sijie-ni-0214 <[email protected]>
This commit is contained in:
sijie-ni-0214
2026-09-29 01:47:36 +08:00
committed by GitHub
parent ad4887b38c
commit 24db48a7e6
41 changed files with 5652 additions and 1024 deletions
Download Patch File Download Diff File Expand all files Collapse all files
configs/milvus.yaml
+5
View File
@@ -614,6 +614,11 @@ queryCoord:
# Set to 0 to disable the penalty period.
resourceExhaustionPenaltyDuration: 30
resourceExhaustionCleanupInterval: 10 # Interval (in seconds) for cleaning up expired resource exhaustion marks on query nodes.
autoscale:
precheckEnabled: false # Whether QueryCoord performs load resource precheck before committing load configuration changes and publishes demand counters after successful broadcasts. With tiered eviction enabled, only demand estimation is performed; capacity checks are skipped.
enabled: false # Whether QueryCoord load resource precheck can use autoscale upper bounds to admit a load request when current QueryNode capacity is insufficient. This is effective only when precheckEnabled is true and tiered eviction is disabled.
maxMemoryLimit: 0 # Maximum global QueryNode memory capacity after autoscale, unit: GiB (1 GiB = 1024^3 bytes). 0 means autoscale cannot add memory capacity.
maxDiskLimit: 0 # Maximum sum of Worker QueryNode LOCAL_STORAGE_SIZE after autoscale, unit: GiB (1 GiB = 1024^3 bytes). 0 means autoscale cannot add disk capacity.
cleanExcludeSegmentInterval: 60 # the time duration of clean pipeline exclude segment which used for filter invalid data, in seconds
ip: # TCP/IP address of queryCoord. If not specified, use the first unicastable address
port: 19531 # TCP port of queryCoord
deployments/monitor/grafana/README.md
+1
View File
@@ -74,6 +74,7 @@ We recommend reading [Milvus monitoring framework overview](monitor_overview.md)
| Parent Load Task | The number of parent load tasks. | ``` sum(milvus_querycoord_parent_task_num{app_kubernetes_io_instance=~"$instance", app_kubernetes_io_name="$app_name", namespace="$namespace"}) by (app_kubernetes_io_instance) ``` | `milvus_querycoord_parent_task_num` | The number of sub load tasks. <br/> Each load request corresponds to a parent task in the task queue. |
| Sub-Load Task Latency | The average latency and the 99th percentile of the latency of a sub load task within the past two minutes. | p99: <br/> ``` histogram_quantile(0.99, sum by (le) (rate(milvus_querycoord_child_task_latency_bucket{app_kubernetes_io_instance=~"$instance", app_kubernetes_io_name="$app_name", namespace="$namespace"}[2m]))) ``` <br/> avg: <br/> ``` sum(increase(milvus_querycoord_child_task_latency_sum{app_kubernetes_io_instance=~"$instance", app_kubernetes_io_name="$app_name", namespace="$namespace"}[2m])) / sum(increase(milvus_querycoord_child_task_latency_count{app_kubernetes_io_instance=~"$instance", app_kubernetes_io_name="$app_name", namespace="$namespace"}[2m])) namespace"}[2m]))) ``` | `milvus_querycoord_child_task_latency` | The latency to complete a sub load task. |
| Query Node Num | The number of query nodes managed by query coord. | ``` sum(milvus_querycoord_querynode_num{app_kubernetes_io_instance=~"$instance", app_kubernetes_io_name="$app_name", namespace="$namespace"}) by (app_kubernetes_io_instance) ``` | `milvus_querycoord_querynode_num` | The number of query nodes managed by query coord. |
| Load Resource Demand | Cumulative estimated memory and disk demand from successfully broadcast load configuration changes, grouped by resource group. | memory: <br/> ``` sum(milvus_querycoord_load_demand_memory_bytes{app_kubernetes_io_instance=~"$instance", app_kubernetes_io_name="$app_name", namespace="$namespace"}) by (rg) / 1024 / 1024 ``` <br/> disk: <br/> ``` sum(milvus_querycoord_load_demand_disk_bytes{app_kubernetes_io_instance=~"$instance", app_kubernetes_io_name="$app_name", namespace="$namespace"}) by (rg) / 1024 / 1024 ``` | `milvus_querycoord_load_demand_memory_bytes`, `milvus_querycoord_load_demand_disk_bytes` | Cumulative estimated load resource demand; external consumers calculate and consume non-overlapping counter deltas. |
</details>
deployments/monitor/grafana/milvus-dashboard.json
+215 -1
View File
@@ -7003,6 +7003,220 @@
"yaxis": {
"align": false
}
},
{
"aliasColors": {},
"bars": false,
"dashLength": 10,
"dashes": false,
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"description": "Cumulative estimated memory demand from successfully broadcast load configuration changes, grouped by resource group, with total.",
"fill": 1,
"fillGradient": 0,
"gridPos": {
"h": 8,
"w": 8,
"x": 0,
"y": 61
},
"hiddenSeries": false,
"id": 123548,
"legend": {
"avg": false,
"current": false,
"max": false,
"min": false,
"show": true,
"total": false,
"values": false
},
"lines": true,
"linewidth": 1,
"nullPointMode": "null",
"options": {
"alertThreshold": true
},
"percentage": false,
"pluginVersion": "10.2.0",
"pointradius": 2,
"points": false,
"renderer": "flot",
"seriesOverrides": [],
"spaceLength": 10,
"stack": false,
"steppedLine": false,
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"editorMode": "code",
"exemplar": true,
"expr": "sum(milvus_querycoord_load_demand_memory_bytes{app_kubernetes_io_instance=~\"$instance\", app_kubernetes_io_name=\"$app_name\", namespace=\"$namespace\"}) by (rg) / 1024 / 1024",
"interval": "",
"legendFormat": "{{rg}}",
"queryType": "randomWalk",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"editorMode": "code",
"exemplar": true,
"expr": "sum(milvus_querycoord_load_demand_memory_bytes{app_kubernetes_io_instance=~\"$instance\", app_kubernetes_io_name=\"$app_name\", namespace=\"$namespace\"}) / 1024 / 1024",
"interval": "",
"legendFormat": "total",
"queryType": "randomWalk",
"range": true,
"refId": "B"
}
],
"thresholds": [],
"timeRegions": [],
"title": "Cumulative Load Demand Memory (MB)",
"tooltip": {
"shared": true,
"sort": 0,
"value_type": "individual"
},
"type": "graph",
"xaxis": {
"mode": "time",
"show": true,
"values": []
},
"yaxes": [
{
"$$hashKey": "object:161",
"format": "short",
"logBase": 1,
"min": "0",
"show": true
},
{
"$$hashKey": "object:162",
"format": "short",
"logBase": 1,
"show": true
}
],
"yaxis": {
"align": false
}
},
{
"aliasColors": {},
"bars": false,
"dashLength": 10,
"dashes": false,
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"description": "Cumulative estimated disk demand from successfully broadcast load configuration changes, grouped by resource group, with total.",
"fill": 1,
"fillGradient": 0,
"gridPos": {
"h": 8,
"w": 8,
"x": 8,
"y": 61
},
"hiddenSeries": false,
"id": 123549,
"legend": {
"avg": false,
"current": false,
"max": false,
"min": false,
"show": true,
"total": false,
"values": false
},
"lines": true,
"linewidth": 1,
"nullPointMode": "null",
"options": {
"alertThreshold": true
},
"percentage": false,
"pluginVersion": "10.2.0",
"pointradius": 2,
"points": false,
"renderer": "flot",
"seriesOverrides": [],
"spaceLength": 10,
"stack": false,
"steppedLine": false,
"targets": [
{
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"editorMode": "code",
"exemplar": true,
"expr": "sum(milvus_querycoord_load_demand_disk_bytes{app_kubernetes_io_instance=~\"$instance\", app_kubernetes_io_name=\"$app_name\", namespace=\"$namespace\"}) by (rg) / 1024 / 1024",
"interval": "",
"legendFormat": "{{rg}}",
"queryType": "randomWalk",
"range": true,
"refId": "A"
},
{
"datasource": {
"type": "prometheus",
"uid": "${datasource}"
},
"editorMode": "code",
"exemplar": true,
"expr": "sum(milvus_querycoord_load_demand_disk_bytes{app_kubernetes_io_instance=~\"$instance\", app_kubernetes_io_name=\"$app_name\", namespace=\"$namespace\"}) / 1024 / 1024",
"interval": "",
"legendFormat": "total",
"queryType": "randomWalk",
"range": true,
"refId": "B"
}
],
"thresholds": [],
"timeRegions": [],
"title": "Cumulative Load Demand Disk (MB)",
"tooltip": {
"shared": true,
"sort": 0,
"value_type": "individual"
},
"type": "graph",
"xaxis": {
"mode": "time",
"show": true,
"values": []
},
"yaxes": [
{
"$$hashKey": "object:163",
"format": "short",
"logBase": 1,
"min": "0",
"show": true
},
{
"$$hashKey": "object:164",
"format": "short",
"logBase": 1,
"show": true
}
],
"yaxis": {
"align": false
}
}
],
"title": "Query Coordinator",
@@ -15608,4 +15822,4 @@
"uid": "uLf5cJ3Gaa",
"version": 9,
"weekStart": ""
}
}
docs/design-docs/design_docs/20260721-querycoord-autoscale-load-resource-precheck.md
+392
View File
@@ -0,0 +1,392 @@
# MEP: Autoscale-Aware Load Resource Admission in QueryCoord
- **Created:** 2026-07-21
- **Updated:** 2026-09-14
- **Author(s):** @sijie-ni-0214
- **Status:** Under Review
- **Component:** QueryCoord / QueryNode / Proxy
- **Related Issue:** [https://github.com/milvus-io/milvus/issues/51217](https://github.com/milvus-io/milvus/issues/51217)
- **Related PR:** [https://github.com/milvus-io/milvus/pull/51218](https://github.com/milvus-io/milvus/pull/51218)
## Summary
QueryCoord evaluates the incremental memory and local-disk demand of
`LoadCollection` and `LoadPartitions` before broadcasting the load
configuration. It admits the request when each resource group has sufficient
capacity, or when the configured global autoscale limits can cover the total
shortage. A resource-insufficient response includes a suggested scale-out
percentage. With tiered eviction enabled, QueryCoord estimates demand without
checking physical availability, because existing caches can be evicted.
After an accepted request's broadcast succeeds, QueryCoord publishes its
incremental demand through per-resource-group Prometheus counters. An external
control plane combines these increments with QueryNode resource waterlines to
make scaling decisions. The feature is optional and disabled by default;
Worker QueryNode admission remains the final node-level resource guard.
## Motivation
Loading a collection introduces demand before QueryNode resource waterlines
reflect it. For example, a collection can be unloaded while new data is
written, then loaded with additional partitions or replicas. A control plane
watching only current memory and disk usage sees little pressure until segment
loading begins. If capacity is insufficient, users encounter repeated load
failures or slow convergence after submitting the request.
QueryCoord knows the requested partitions, fields, indexes, replica placement,
and segment metadata. It can estimate the demand at the point where the load
configuration is accepted. QueryNodes supply memory capacity and usage together
with local-storage quota and directory usage. The external control plane owns
instance selection, node provisioning, resource-group placement, cooldown,
and product limits.
The design gives the caller an early capacity decision and gives the control
plane a demand signal in bytes. This keeps admission close to the load
configuration while leaving infrastructure policy with the control plane.
## Goals and Non-Goals
The goals are to:
- estimate incremental Worker QN memory and local-disk demand for a load
configuration;
- evaluate demand independently for each affected resource group;
- support admission within configured global autoscale upper bounds;
- expose a scale-out suggestion and accepted-demand metrics;
- reuse QN resource-estimation formulas and tolerate unavailable estimation
dependencies.
Node provisioning, scale-from-zero, per-segment placement, and transactional
resource reservations are outside the scope. The estimate covers Worker QN
segment resources; SQN delegator resources, including Bloom and BM25 statistics
loaded through their own paths, require separate accounting.
## Design Details
### Load request flow
The precheck runs under the collection broadcast lock after replica assignment
and load-configuration generation, and before broadcast. Different collections
can execute concurrently.
```mermaid
flowchart TD
Load["LoadCollection / LoadPartitions"] --> Config["Resolve configuration and replica placement"]
Config --> Check["Estimate incremental demand and evaluate capacity"]
Check -- "Current capacity or autoscale headroom fits" --> Broadcast["Broadcast load configuration"]
Check -- "Tiered eviction: retain demand, skip capacity check" --> Broadcast
Check -- "Dependency error: fail open" --> Broadcast
Check -- "Insufficient resources" --> Reject["Return scale-out suggestion"]
Broadcast -- "Success with accepted demand" --> Counter["Add per-RG demand counters"]
Broadcast -- Success --> Worker["Worker QNs load segments"]
Counter --> Control["External scaling control plane"]
Worker --> Usage["QueryNode resource waterlines"]
Usage --> Control
```
`LoadCollection` uses the collection's partitions. `LoadPartitions` combines
the requested partitions with those already loaded. Each request supplies the
current and expected configuration to the estimator. Refresh requests use the
refresh path, and an unchanged configuration completes before precheck.
Replica assignment checks node counts for an initial collection load and for
partition loads. Failure returns `ErrResourceGroupNodeNotEnough`; the
resource precheck's byte-based suggestion applies after this placement gate.
When precheck is disabled, the request proceeds to broadcast without a demand
estimate. When enabled, it returns an accepted estimate, a resource-shortage
error, or no estimate after a metadata or estimation failure. Capacity-metrics
failures preserve successfully estimated demand while allowing the load to
proceed. The caller records demand only after the broadcast returns success.
### Resource estimation
QueryCoord uses `GetRecoveryInfoV2` to select target segment IDs, then fetches
complete `SegmentInfo` metadata through `GetSegmentInfo` and indexes through
`GetIndexInfo`. The full-segment broker batches up to 1000 IDs and reconstructs
compressed binlog paths; index requests batch up to 1024 IDs.
`ErrIndexNotFound` is an accepted absence of index metadata.
The shared estimator in `internal/util/segcore/loadresource` computes final
memory and disk footprint from schema, indexes, raw data, regular stats, delta
logs, and JSON/Text statistics. Index estimates use the C++ metadata-based
resource estimator without reading index files. Raw data follows mmap placement
and type-specific overhead rules; Text estimates include index-file bytes and
a separate word-aligned validity bitmap.
Internal Storage V3 segments can supply raw column-group sizes through
`Statistics.load_resource.column_groups` and delete memory through
`Statistics.delta_binlog_size`. A column group records its ID, member fields,
and uncompressed memory total. When raw binlog arrays are present, the estimator
uses the supplied raw/delta arrays. Otherwise, an internal V3 segment with a
manifest and a non-nil load-resource summary uses pathless raw/delta metadata
derived from that summary. This adaptation requires no manifest read and
leaves the load message intact. JSON/Text statistics use their existing
metadata maps.
The final estimator uses a positive binlog `MemorySize`, falling back to
`LogSize` otherwise; the loading estimator uses explicit `MemorySize`.
JSON/Text estimates use `MemorySize` in both paths. A V3 synthetic delta with
positive memory and zero serialized size contributes 1x final memory and 2x
loading memory. An index that supplies raw data can suppress its corresponding
field-binlog entry according to the index and raw-data preference settings.
With tiered eviction enabled, final footprint includes only the configured
cache share of evictable resources:
```text
final_memory = non_evictable_memory + ceil(evictable_memory * memory_cache_ratio)
final_disk = non_evictable_disk + ceil(evictable_disk * disk_cache_ratio)
```
With tiered eviction disabled, the complete evictable footprint is included.
Zero cache ratios are valid. QN logical accounting uses the same final
estimator and caches its result before compacting runtime load metadata.
QN `requestResource` uses the separate loading-peak estimator, which also
accounts for transient allocations.
QueryCoord filters each configuration's metadata by partitions and load fields.
All indexes belonging to a loaded field are retained, including multiple JSON
path indexes on the same field. System fields are retained. A physical column
group is kept in full when any member is selected; its size is not split between
children. Delta applies to the whole segment. The estimator receives the full
schema and ignores dropped fields when classifying final resource usage.
Worker QN admission consumes its load-request metadata directly, so shared
formulas do not imply identical input selection.
Metadata and segment/index estimates execute sequentially within a precheck.
QueryCoord invokes CGo synchronously; QN uses its dynamic pool. The estimate is
best-effort: V3 summaries have no completeness or version check and do not
describe every runtime allocation or child-manifest delete source. Raw-vector
placement uses the global mmap option, so per-field vector overrides can
differ from actual loading.
### Incremental demand
Let `E` be the expected configuration's final footprint and `C` the footprint
represented by the current configuration. Both contain memory and disk bytes.
The current estimate uses current partitions within the fetched target segment
set, current fields/indexes, and the stored schema when available. It represents
configured demand rather than measured replica residency.
For each expected resource group:
```text
kept(rg) = min(current_replicas(rg), expected_replicas(rg))
incoming(rg) = max(expected_replicas(rg) - current_replicas(rg), 0)
required(rg) = kept(rg) * max(E - C, 0) + incoming(rg) * E
```
The formula is applied independently to memory and disk. New replicas
contribute the full expected footprint; retained replicas contribute positive
increases. Reductions do not offset positive demand in another dimension or
RG. Zero-demand groups are omitted.
For example, increasing a replica's footprint from 80 MiB to 120 MiB while
adding a second replica in the same RG requires `40 + 120 = 160 MiB`.
### Capacity and admission
With tiered eviction enabled, QueryCoord returns the estimated demand without
collecting QN capacity metrics or checking node availability and global
autoscale limits. Existing evictable caches make current usage unsuitable for
this admission decision. Successful broadcasts still add the estimated demand
to the counters; Worker QN enforces its loading-stage resource guard.
The following capacity evaluation applies when tiered eviction is disabled.
QueryCoord collects one request-local set of `system_info` samples for the
non-stopping nodes in affected RGs. Autoscale upper-bound evaluation also
includes all online, non-stopping Worker QNs for global capacity, excluding
QueryNodes embedded in StreamingNodes. Node IDs are
deduplicated; collection uses at most 16 concurrent calls and one attempt per
node, with a shared ten-second timeout for the collection round. Any RPC,
status, parsing, empty eligible-node set, or required-capacity failure
invalidates the capacity evaluation while preserving successfully estimated demand.
Each Worker QN reports its configured `LOCAL_STORAGE_SIZE` and cached
local-storage directory usage in `system_info`, both in bytes. Missing or invalid
local-storage data invalidates the capacity evaluation. Filesystem-wide disk
metrics remain available for monitoring but do not enter this precheck.
Worker QNs sharing one filesystem need per-QN `LOCAL_STORAGE_SIZE` quotas
consistent with that shared capacity. The precheck follows each Worker's
configured limit and does not reserve or coordinate filesystem free space
between Workers.
Let `hm` be `queryNodeMemoryHighWaterLevel` and `hd` be the parsed
`maxDiskUsagePercentage` ratio:
```text
node_memory_capacity = physical_memory
node_disk_capacity = LOCAL_STORAGE_SIZE
node_memory_available = node_memory_capacity * hm - memory_usage
node_disk_available = max(node_disk_capacity * hd - local_directory_usage, 0)
rg_available = sum(node_available in RG)
rg_usable_capacity.memory = sum(node_memory_capacity * hm in RG)
rg_usable_capacity.disk = sum(node_disk_capacity * hd in RG)
shortage(rg) = max(required(rg) - rg_available, 0)
total_shortage = sum(shortage(rg))
```
Memory availability can be negative, preserving existing overcommit in the
shortage calculation. Disk availability is clamped to zero per node. Memory
and disk must both fit, and free capacity in another RG does not offset a
group's shortage.
If a shortage exists and autoscale admission is enabled, the request may use
future usable capacity within global limits:
```text
remaining_capacity.memory = max(maxMemoryLimit - global_physical_memory_capacity, 0)
remaining_capacity.disk = max(maxDiskLimit - global_local_storage_capacity, 0)
usable_headroom.memory = remaining_capacity.memory * hm
usable_headroom.disk = remaining_capacity.disk * hd
admit if total_shortage.memory <= usable_headroom.memory
and total_shortage.disk <= usable_headroom.disk
```
Global capacity includes online, non-stopping Worker QNs. Memory sums their
physical memory capacity; disk sums their `LOCAL_STORAGE_SIZE` quotas. The
limits express a cluster-wide upper bound; the control plane supplies the
actual nodes and assigns them to RGs.
For a rejected request, the suggestion uses threshold-adjusted RG capacity:
```text
dimension_percent = ceil(100 * shortage / rg_usable_capacity)
suggested_expand_percent = max(dimension_percent across RGs and dimensions)
```
A positive shortage with zero capacity yields 100 percent. The result is a
capacity hint, not an exact instance shape or node count.
Capacity samples are observations, not reservations. Concurrent collection
loads can see the same availability or headroom, and local directory usage
can lag accepted demand. RG aggregate admission also does not guarantee that an
individual segment fits a particular node. Worker QN's loading guard remains
responsible for those node-level decisions.
### Demand metrics
QueryCoord exposes two counters, each with the single application label `rg`:
| Metric | Meaning |
|---|---|
| `milvus_querycoord_load_demand_memory_bytes` | Cumulative accepted incremental memory demand whose load-config broadcast returned success. |
| `milvus_querycoord_load_demand_disk_bytes` | Cumulative accepted incremental local-disk demand whose load-config broadcast returned success. |
A counter increases only when precheck produces positive accepted demand and
the broadcast succeeds. Disabled checks, unchanged configurations, zero demand,
metadata or estimation failures, resource rejections, and failed broadcasts
produce no increment. Fail-open requests caused by unavailable capacity metrics
retain their estimated demand and increment counters after successful broadcast.
Load completion and release do not decrement the counters.
Default and named RG series are initialized at zero during registration and
RG setup; RG deletion removes the corresponding series. The dashboard
displays cumulative demand per RG and cluster totals in MiB.
An external consumer keeps a baseline per source series and consumes
non-overlapping increments, handling restarts, leader changes, and RG
recreation. It combines demand with QN waterlines and its provisioning state
to avoid acting repeatedly on the same demand or counting demand already
reflected in usage. Overlapping rolling-window queries are not a one-time
consumption protocol.
Counters are best-effort telemetry. A process failure between broadcast,
recording, and scraping can lose an observation; the metric is not a durable
event log or a ledger of outstanding loads.
### Failure and response semantics
Metadata lookup, estimation, and QN metrics errors are logged and fail open.
Metadata or estimation failures leave no demand to publish. QN metrics errors,
collection timeouts, and empty eligible-node sets preserve the already computed
demand for publication after a successful broadcast. A completed resource
evaluation that finds insufficient capacity returns
`ErrServiceResourceInsufficient` with:
```text
Retriable = false
ExtraInfo["suggested_expand_percent"] = "<integer percentage>"
```
Proxy forwards the non-OK status and its extra fields intact. Transport
errors and nil responses remain task errors. A resource rejection remains
an application failure even when the transport returns no Go error.
`LoadCollection` post-execution cache work is skipped for a non-OK status.
## Configuration and Compatibility
```yaml
queryCoord:
autoscale:
precheckEnabled: false
enabled: false
maxMemoryLimit: 0
maxDiskLimit: 0
```
All four settings are refreshable:
| Key under `queryCoord.autoscale` | Meaning |
|---|---|
| `precheckEnabled` | Enable resource precheck and accepted-demand recording. |
| `enabled` | Allow admission using global autoscale upper bounds. |
| `maxMemoryLimit` | Maximum global QN physical memory capacity after scaling, in GiB. |
| `maxDiskLimit` | Maximum sum of Worker QN `LOCAL_STORAGE_SIZE` quotas after scaling, in GiB. |
With precheck disabled, normal load behavior applies. With precheck enabled
and tiered eviction disabled, admission uses current capacity unless autoscale
upper-bound admission is enabled. A zero global limit
cannot cover a positive shortage. Enabling both switches permits admission
within configured limits when tiered eviction is disabled, but does not itself
provision resources. With tiered eviction enabled, precheck estimates demand
only, regardless of the autoscale switch and capacity limits.
Configured capacity limits use GiB: `1 GiB = 1024^3 bytes`. QN `GetMetrics`
local-storage quota and directory usage, memory metrics, load demand counters,
and all admission calculations use bytes. Filesystem-wide disk metrics use
decimal GB (`1 GB = 1e9 bytes`) for monitoring and do not enter admission.
Deployments should keep `queryNodeMemoryHighWaterLevel` no higher than the
non-tiered Worker guard's `overloadedMemoryThresholdPercentage`. QueryCoord
does not enforce this relationship. The Worker guard remains authoritative for
individual loads, including its in-flight resource reservations.
The public load request shape is unchanged. The scale-out suggestion uses
`commonpb.Status.ExtraInfo`; clients may consume it alongside the error reason.
Estimate accuracy depends on supplied metadata and the estimator's local QN
configuration, including mmap and tiered settings.
## Test Plan
Unit tests cover the following behavior:
- incremental demand for initial loads, partition/field/index changes,
retained replicas, and incoming replicas;
- complete segment metadata lookup, V3 column-group summaries, and
group-preserving field selection;
- current-capacity and global-limit decisions using Worker QN local-storage
quotas independently of filesystem-wide disk metrics, including nodes that
share a filesystem, negative memory availability, percentage rounding, and
zero-capacity handling;
- shared final/loading estimates, mmap and tiered settings, JSON/Text
statistics, and Text validity bitmap memory;
- bounded metrics collection, snapshot reuse, dependency failures, and
preservation of resource-insufficient statuses;
- counter initialization, RG lifecycle, additive recording, and separation
of accepted-demand evaluation from publication;
- Proxy status propagation and QN logical-estimate caching.
Integration validation should cover the complete load request through
DataCoord metadata lookup, QueryCoord broadcast, and Worker QN loading,
including V3 loads after metadata recovery. It should verify that only
successfully broadcast, evaluated demand reaches the counters and that
metadata or estimation failures permit loading without a metric increment,
while capacity-metrics failures preserve demand for successful broadcasts. Control-plane
validation should exercise counter resets and demand already reflected in
waterlines. Broker-mock tests alone do not establish these end-to-end results.
internal/core/src/segcore/load_index_c.cpp
+108 -78
View File
@@ -58,6 +58,87 @@
#include "storage/RemoteChunkManagerSingleton.h"
#include "storage/Util.h"
namespace {
void
FillLoadIndexInfoFromProto(const milvus::proto::cgo::LoadIndexInfo& info_proto,
milvus::segcore::LoadIndexInfo* load_index_info) {
load_index_info->collection_id = info_proto.collectionid();
load_index_info->partition_id = info_proto.partitionid();
load_index_info->segment_id = info_proto.segmentid();
load_index_info->field_id = info_proto.field().fieldid();
load_index_info->field_type =
static_cast<milvus::DataType>(info_proto.field().data_type());
load_index_info->element_type =
static_cast<milvus::DataType>(info_proto.field().element_type());
load_index_info->enable_mmap = info_proto.enable_mmap();
load_index_info->index_id = info_proto.indexid();
load_index_info->index_build_id = info_proto.index_buildid();
load_index_info->index_version = info_proto.index_version();
load_index_info->index_store_path_version =
info_proto.index_store_path_version();
for (const auto& [k, v] : info_proto.index_params()) {
load_index_info->index_params[k] = v;
}
load_index_info->index_files.assign(info_proto.index_files().begin(),
info_proto.index_files().end());
load_index_info->uri = info_proto.uri();
load_index_info->index_engine_version = info_proto.index_engine_version();
auto scalar_version = info_proto.current_scalar_index_version();
if (scalar_version > 0) {
load_index_info
->index_params[milvus::index::SCALAR_INDEX_ENGINE_VERSION] =
std::to_string(scalar_version);
}
load_index_info->schema = info_proto.field();
load_index_info->index_size = info_proto.index_file_size();
load_index_info->num_rows = info_proto.num_rows();
auto field_schema = milvus::FieldMeta::ParseFrom(load_index_info->schema);
size_t dim =
IsVectorDataType(field_schema.get_data_type()) &&
!IsSparseFloatVectorDataType(field_schema.get_data_type())
? field_schema.get_dim()
: 1;
load_index_info->dim = dim;
auto warmup_it = load_index_info->index_params.find("warmup");
if (warmup_it != load_index_info->index_params.end()) {
load_index_info->warmup_policy = warmup_it->second;
LOG_INFO("Index warmup_policy extracted from index_params: {}",
load_index_info->warmup_policy);
} else {
load_index_info->warmup_policy = "";
LOG_INFO("No warmup key in index_params, warmup_policy will be empty");
}
}
LoadResourceRequest
EstimateLoadIndexResourceFromLoadInfo(
const milvus::segcore::LoadIndexInfo* load_index_info) {
auto field_type = load_index_info->field_type;
auto element_type = load_index_info->element_type;
auto& index_params = load_index_info->index_params;
bool find_index_type = index_params.count("index_type") > 0 ? true : false;
if (!(find_index_type == true)) {
ThrowInfo(milvus::ErrorCode::DataFormatBroken,
"Can't find index type in index_params");
}
// Keep admission estimation metadata-only: exact scalar V3 directory
// inspection belongs to SealedIndexTranslator, where the result is used
// for the actual MCL loading reservation.
return milvus::index::IndexFactory::GetInstance().IndexLoadResource(
field_type,
element_type,
load_index_info->index_engine_version,
load_index_info->index_size,
index_params,
load_index_info->enable_mmap,
load_index_info->num_rows,
load_index_info->dim);
}
} // namespace
bool
IsLoadWithDisk(const char* index_type, int index_engine_version) {
SCOPE_CGO_CALL_METRIC();
@@ -103,31 +184,8 @@ EstimateLoadIndexResource(CLoadIndexInfo c_load_index_info,
try {
auto load_index_info =
(milvus::segcore::LoadIndexInfo*)c_load_index_info;
auto field_type = load_index_info->field_type;
auto element_type = load_index_info->element_type;
auto& index_params = load_index_info->index_params;
bool find_index_type =
index_params.count("index_type") > 0 ? true : false;
if (!(find_index_type == true)) {
ThrowInfo(milvus::ErrorCode::DataFormatBroken,
"Can't find index type in index_params");
}
// Segment Loader calls this API while deciding whether a segment may
// start loading. Keep that admission path metadata-only: exact scalar
// V3 directory inspection belongs to SealedIndexTranslator, where the
// result is used for the actual MCL loading reservation.
LoadResourceRequest request =
milvus::index::IndexFactory::GetInstance().IndexLoadResource(
field_type,
element_type,
load_index_info->index_engine_version,
load_index_info->index_size,
index_params,
load_index_info->enable_mmap,
load_index_info->num_rows,
load_index_info->dim);
*c_load_resource_request = request;
*c_load_resource_request =
EstimateLoadIndexResourceFromLoadInfo(load_index_info);
return milvus::SuccessCStatus();
}
// Estimation failure must reach the Go caller as an error: a swallowed
@@ -136,6 +194,30 @@ EstimateLoadIndexResource(CLoadIndexInfo c_load_index_info,
CGO_CATCH_AND_RETURN_CSTATUS
}
CStatus
EstimateLoadIndexResourceFromSerializedInfo(
const uint8_t* serialized_load_index_info,
const uint64_t len,
LoadResourceRequest* load_resource_request) {
SCOPE_CGO_CALL_METRIC();
try {
AssertInfo(load_resource_request != nullptr,
"load resource request is null");
auto info_proto = milvus::proto::cgo::LoadIndexInfo();
auto parsed =
info_proto.ParseFromArray(serialized_load_index_info, len);
AssertInfo(parsed, "failed to parse load index info");
auto load_index_info = milvus::segcore::LoadIndexInfo();
FillLoadIndexInfoFromProto(info_proto, &load_index_info);
*load_resource_request =
EstimateLoadIndexResourceFromLoadInfo(&load_index_info);
return milvus::SuccessCStatus();
}
CGO_CATCH_AND_RETURN_CSTATUS
}
bool
TryReserveLoadingResourceWithTimeout(CResourceUsage size,
int64_t millisecond_timeout) {
@@ -249,59 +331,7 @@ FinishLoadIndexInfo(CLoadIndexInfo c_load_index_info,
static_cast<milvus::segcore::LoadIndexInfo*>(c_load_index_info);
// TODO: keep this since LoadIndexInfo is used by SegmentSealed.
{
load_index_info->collection_id = info_proto->collectionid();
load_index_info->partition_id = info_proto->partitionid();
load_index_info->segment_id = info_proto->segmentid();
load_index_info->field_id = info_proto->field().fieldid();
load_index_info->field_type =
static_cast<milvus::DataType>(info_proto->field().data_type());
load_index_info->element_type = static_cast<milvus::DataType>(
info_proto->field().element_type());
load_index_info->enable_mmap = info_proto->enable_mmap();
load_index_info->index_id = info_proto->indexid();
load_index_info->index_build_id = info_proto->index_buildid();
load_index_info->index_version = info_proto->index_version();
load_index_info->index_store_path_version =
info_proto->index_store_path_version();
for (const auto& [k, v] : info_proto->index_params()) {
load_index_info->index_params[k] = v;
}
load_index_info->index_files.assign(
info_proto->index_files().begin(),
info_proto->index_files().end());
load_index_info->uri = info_proto->uri();
load_index_info->index_engine_version =
info_proto->index_engine_version();
// Inject scalar index version into index_params for scalar indexes
auto scalar_version = info_proto->current_scalar_index_version();
if (scalar_version > 0) {
load_index_info
->index_params[milvus::index::SCALAR_INDEX_ENGINE_VERSION] =
std::to_string(scalar_version);
}
load_index_info->schema = info_proto->field();
load_index_info->index_size = info_proto->index_file_size();
load_index_info->num_rows = info_proto->num_rows();
auto field_schema =
milvus::FieldMeta::ParseFrom(load_index_info->schema);
size_t dim = IsVectorDataType(field_schema.get_data_type()) &&
!IsSparseFloatVectorDataType(
field_schema.get_data_type())
? field_schema.get_dim()
: 1;
load_index_info->dim = dim;
// Extract warmup_policy from index_params (keep it for Knowhere)
auto warmup_it = load_index_info->index_params.find("warmup");
if (warmup_it != load_index_info->index_params.end()) {
load_index_info->warmup_policy = warmup_it->second;
LOG_INFO("Index warmup_policy extracted from index_params: {}",
load_index_info->warmup_policy);
} else {
LOG_INFO(
"No warmup key in index_params, warmup_policy will be "
"empty");
}
FillLoadIndexInfoFromProto(*info_proto, load_index_info);
auto remote_chunk_manager =
milvus::storage::RemoteChunkManagerSingleton::GetInstance()
.GetRemoteChunkManager();
internal/core/src/segcore/load_index_c.h
+6
View File
@@ -36,6 +36,12 @@ CStatus
EstimateLoadIndexResource(CLoadIndexInfo c_load_index_info,
LoadResourceRequest* c_load_resource_request);
CStatus
EstimateLoadIndexResourceFromSerializedInfo(
const uint8_t* serialized_load_index_info,
const uint64_t len,
LoadResourceRequest* load_resource_request);
bool
TryReserveLoadingResourceWithTimeout(CResourceUsage size,
int64_t millisecond_timeout);
internal/proxy/task.go
+12 -2
View File
@@ -3146,13 +3146,20 @@ func (t *loadCollectionTask) Execute(ctx context.Context) (err error) {
mlog.FieldSchema(request.Schema),
mlog.Int32("priority", int32(request.GetPriority())))
t.result, err = t.mixCoord.LoadCollection(ctx, request)
if err = merr.CheckRPCCall(t.result, err); err != nil {
if err != nil {
return merr.Wrap(err, "call query coordinator LoadCollection")
}
if t.result == nil {
return merr.CheckRPCCall(t.result, nil)
}
return nil
}
func (t *loadCollectionTask) PostExecute(ctx context.Context) error {
if t.result != nil && !merr.Ok(t.result) {
return nil
}
collID, err := t.GetMetaCache().GetCollectionID(ctx, t.GetDbName(), t.CollectionName)
mlog.Debug(ctx, "loadCollectionTask PostExecute",
mlog.String("role", typeutil.ProxyRole),
@@ -3418,9 +3425,12 @@ func (t *loadPartitionsTask) Execute(ctx context.Context) error {
mlog.FieldSchema(request.Schema),
mlog.Int32("priority", int32(request.GetPriority())))
t.result, err = t.mixCoord.LoadPartitions(ctx, request)
if err = merr.CheckRPCCall(t.result, err); err != nil {
if err != nil {
return err
}
if t.result == nil {
return merr.CheckRPCCall(t.result, nil)
}
return nil
}
internal/proxy/task_test.go
+18
View File
@@ -3885,6 +3885,24 @@ func TestLoadCollectionTaskExecuteTextRequiresStorageV3(t *testing.T) {
assert.Contains(t, err.Error(), "TEXT field requires StorageV3")
}
func TestLoadCollectionTaskPostExecuteSkipsFailedStatus(t *testing.T) {
task := &loadCollectionTask{
LoadCollectionRequest: &milvuspb.LoadCollectionRequest{
DbName: "db",
CollectionName: "collection",
},
result: &commonpb.Status{
Code: merr.Code(merr.ErrServiceResourceInsufficient),
Reason: "insufficient resource",
ExtraInfo: map[string]string{
"suggested_expand_percent": "10",
},
},
}
assert.NoError(t, task.PostExecute(context.Background()))
}
func Test_loadPartitionTask_Execute(t *testing.T) {
qc := NewMixCoordMock()
internal/querycoordv2/autoscale/resource.go
+264
View File
@@ -0,0 +1,264 @@
// 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 autoscale
import (
"context"
"github.com/samber/lo"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/util/segcore/loadresource"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type ResourceUsage struct {
MemoryBytes int64
DiskBytes int64
}
type EstimateOptions struct {
MmapVectorField bool
MmapScalarField bool
MmapJSONStats bool
TieredEvictionEnabled bool
TieredMemoryCacheRatio float64
TieredDiskCacheRatio float64
DeltaDataExpansionFactor float64
JSONKeyStatsExpansionFactor float64
TextIndexExpansionFactor float64
PreferFieldDataWhenIndexHasRawData bool
}
func DefaultEstimateOptions() EstimateOptions {
queryNodeCfg := &paramtable.Get().QueryNodeCfg
return EstimateOptions{
MmapVectorField: queryNodeCfg.MmapVectorField.GetAsBool(),
MmapScalarField: queryNodeCfg.MmapScalarField.GetAsBool(),
MmapJSONStats: queryNodeCfg.MmapJSONStats.GetAsBool(),
TieredEvictionEnabled: queryNodeCfg.TieredEvictionEnabled.GetAsBool(),
TieredMemoryCacheRatio: queryNodeCfg.TieredEvictableMemoryCacheRatio.GetAsFloat(),
TieredDiskCacheRatio: queryNodeCfg.TieredEvictableDiskCacheRatio.GetAsFloat(),
DeltaDataExpansionFactor: queryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
JSONKeyStatsExpansionFactor: queryNodeCfg.JSONKeyStatsExpansionFactor.GetAsFloat(),
TextIndexExpansionFactor: queryNodeCfg.TextIndexExpansionFactor.GetAsFloat(),
PreferFieldDataWhenIndexHasRawData: queryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool(),
}
}
func EstimateSegmentsLoadResource(ctx context.Context, schema *schemapb.CollectionSchema, segments []*datapb.SegmentInfo, indexes map[int64][]*querypb.FieldIndexInfo, options EstimateOptions) (ResourceUsage, error) {
options = options.normalized()
if schema == nil {
return ResourceUsage{}, merr.WrapErrServiceInternalMsg("collection schema is nil")
}
usage := ResourceUsage{}
for _, segment := range segments {
if err := ctx.Err(); err != nil {
return ResourceUsage{}, err
}
segmentIndexes := indexes[segment.GetID()]
finalUsage, err := loadresource.EstimateSegmentFinalResource(
ctx,
schema,
segmentLoadInfo(segment, segmentIndexes),
options.segmentFinalEstimateOptions(),
nil,
)
if err != nil {
return ResourceUsage{}, merr.Wrapf(err, "failed to estimate load resource for segment %d", segment.GetID())
}
usage.add(ResourceUsage{
MemoryBytes: uint64ToInt64(finalUsage.MemoryBytes),
DiskBytes: uint64ToInt64(finalUsage.DiskBytes),
})
}
if err := ctx.Err(); err != nil {
return ResourceUsage{}, err
}
return usage, nil
}
func EstimateSegmentsLoadResourceForLoadConfig(ctx context.Context, schema *schemapb.CollectionSchema, segments []*datapb.SegmentInfo, indexes map[int64][]*querypb.FieldIndexInfo, loadFields []int64, options EstimateOptions) (ResourceUsage, error) {
return EstimateSegmentsLoadResource(
ctx,
schema,
FilterSegmentsByLoadFields(segments, loadFields),
FilterIndexesByLoadFields(indexes, loadFields),
options,
)
}
func FilterSegmentsByLoadFields(segments []*datapb.SegmentInfo, loadFields []int64) []*datapb.SegmentInfo {
if len(loadFields) == 0 {
return segments
}
loadFieldSet := typeutil.NewSet(loadFields...)
return lo.Map(segments, func(segment *datapb.SegmentInfo, _ int) *datapb.SegmentInfo {
filtered := proto.Clone(segment).(*datapb.SegmentInfo)
filtered.Binlogs = filterFieldBinlogsByLoadFields(filtered.GetBinlogs(), loadFieldSet)
filtered.Statslogs = filterFieldBinlogsByLoadFields(filtered.GetStatslogs(), loadFieldSet)
filtered.Bm25Statslogs = filterFieldBinlogsByLoadFields(filtered.GetBm25Statslogs(), loadFieldSet)
filtered.JsonKeyStats = filterMapByLoadFields(filtered.GetJsonKeyStats(), loadFieldSet)
filtered.TextStatsLogs = filterMapByLoadFields(filtered.GetTextStatsLogs(), loadFieldSet)
filterLoadResourceStatisticsByLoadFields(filtered.GetStats(), loadFieldSet)
return filtered
})
}
func filterLoadResourceStatisticsByLoadFields(stats *datapb.Statistics, loadFields typeutil.Set[int64]) {
resource := stats.GetLoadResource()
if resource == nil {
return
}
resource.ColumnGroups = lo.Filter(resource.GetColumnGroups(), func(group *datapb.ColumnGroupStatistics, _ int) bool {
fieldIDs := group.GetFieldIds()
if len(fieldIDs) == 0 {
return shouldKeepFieldForLoadConfig(group.GetGroupId(), loadFields)
}
return lo.SomeBy(fieldIDs, func(fieldID int64) bool {
return shouldKeepFieldForLoadConfig(fieldID, loadFields)
})
})
}
func filterFieldBinlogsByLoadFields(fieldBinlogs []*datapb.FieldBinlog, loadFields typeutil.Set[int64]) []*datapb.FieldBinlog {
return lo.Filter(fieldBinlogs, func(fieldBinlog *datapb.FieldBinlog, _ int) bool {
childFields := fieldBinlog.GetChildFields()
if len(childFields) == 0 {
return shouldKeepFieldForLoadConfig(fieldBinlog.GetFieldID(), loadFields)
}
for _, fieldID := range childFields {
if shouldKeepFieldForLoadConfig(fieldID, loadFields) {
return true
}
}
return false
})
}
func filterMapByLoadFields[T any](values map[int64]T, loadFields typeutil.Set[int64]) map[int64]T {
if len(values) == 0 {
return values
}
filtered := make(map[int64]T)
for fieldID, value := range values {
if shouldKeepFieldForLoadConfig(fieldID, loadFields) {
filtered[fieldID] = value
}
}
return filtered
}
func shouldKeepFieldForLoadConfig(fieldID int64, loadFields typeutil.Set[int64]) bool {
return common.IsSystemField(fieldID) || loadFields.Contain(fieldID)
}
func FilterIndexesByLoadFields(indexes map[int64][]*querypb.FieldIndexInfo, loadFields []int64) map[int64][]*querypb.FieldIndexInfo {
if len(indexes) == 0 || len(loadFields) == 0 {
return indexes
}
loadFieldSet := typeutil.NewSet(loadFields...)
filtered := make(map[int64][]*querypb.FieldIndexInfo)
for segmentID, segmentIndexes := range indexes {
for _, index := range segmentIndexes {
if !loadFieldSet.Contain(index.GetFieldID()) {
continue
}
filtered[segmentID] = append(filtered[segmentID], index)
}
}
return filtered
}
func (options EstimateOptions) normalized() EstimateOptions {
if options.TieredMemoryCacheRatio < 0 || options.TieredMemoryCacheRatio > 1 {
options.TieredMemoryCacheRatio = 1
}
if options.TieredDiskCacheRatio < 0 || options.TieredDiskCacheRatio > 1 {
options.TieredDiskCacheRatio = 1
}
if options.DeltaDataExpansionFactor <= 0 {
options.DeltaDataExpansionFactor = 1
}
if options.JSONKeyStatsExpansionFactor <= 0 {
options.JSONKeyStatsExpansionFactor = 1
}
if options.TextIndexExpansionFactor <= 0 {
options.TextIndexExpansionFactor = 1
}
return options
}
func (options EstimateOptions) segmentFinalEstimateOptions() loadresource.SegmentFinalEstimateOptions {
return loadresource.SegmentFinalEstimateOptions{
MmapVectorField: options.MmapVectorField,
MmapScalarField: options.MmapScalarField,
MmapJSONStats: options.MmapJSONStats,
DeltaDataExpansionFactor: options.DeltaDataExpansionFactor,
JSONKeyStatsExpansionFactor: options.JSONKeyStatsExpansionFactor,
TextIndexExpansionFactor: options.TextIndexExpansionFactor,
TieredEvictionEnabled: options.TieredEvictionEnabled,
TieredEvictableMemoryCacheRatio: options.TieredMemoryCacheRatio,
TieredEvictableDiskCacheRatio: options.TieredDiskCacheRatio,
PreferFieldDataWhenIndexHasRawData: options.PreferFieldDataWhenIndexHasRawData,
}
}
func segmentLoadInfo(segment *datapb.SegmentInfo, indexes []*querypb.FieldIndexInfo) *querypb.SegmentLoadInfo {
return &querypb.SegmentLoadInfo{
CollectionID: segment.GetCollectionID(),
PartitionID: segment.GetPartitionID(),
SegmentID: segment.GetID(),
NumOfRows: segment.GetNumOfRows(),
BinlogPaths: segment.GetBinlogs(),
Statslogs: segment.GetStatslogs(),
Deltalogs: segment.GetDeltalogs(),
IndexInfos: indexes,
Bm25Logs: segment.GetBm25Statslogs(),
JsonKeyStatsLogs: segment.GetJsonKeyStats(),
TextStatsLogs: segment.GetTextStatsLogs(),
InsertChannel: segment.GetInsertChannel(),
Level: segment.GetLevel(),
StorageVersion: segment.GetStorageVersion(),
IsSorted: segment.GetIsSorted(),
ManifestPath: segment.GetManifestPath(),
CommitTimestamp: segment.GetCommitTimestamp(),
DataVersion: segment.GetDataVersion(),
CompactionFrom: segment.GetCompactionFrom(),
Stats: segment.GetStats(),
}
}
func uint64ToInt64(value uint64) int64 {
const maxInt64 = uint64(1<<63 - 1)
if value > maxInt64 {
return int64(maxInt64)
}
return int64(value)
}
func (usage *ResourceUsage) add(other ResourceUsage) {
usage.MemoryBytes += other.MemoryBytes
usage.DiskBytes += other.DiskBytes
}
internal/querycoordv2/autoscale/resource_test.go
+507
View File
@@ -0,0 +1,507 @@
// 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 autoscale
import (
"context"
"testing"
"github.com/samber/lo"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"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/internal/storage"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
)
func TestEstimateSegmentsLoadResourceReturnsErrorWithoutSchema(t *testing.T) {
usage, err := EstimateSegmentsLoadResource(context.Background(), nil, nil, nil, EstimateOptions{})
assert.ErrorIs(t, err, merr.ErrServiceInternal)
assert.Zero(t, usage)
}
func TestEstimateSegmentsLoadResourceStopsOnCanceledContext(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
usage, err := EstimateSegmentsLoadResource(ctx, testSchema(), []*datapb.SegmentInfo{{ID: 1}}, nil, EstimateOptions{})
require.ErrorIs(t, err, context.Canceled)
assert.Zero(t, usage)
}
func TestEstimateSegmentsLoadResourceUsesSchemaAndMmap(t *testing.T) {
schema := testSchema()
segments := []*datapb.SegmentInfo{
{
ID: 1,
NumOfRows: 1000,
Binlogs: testRawBinlogs(),
Statslogs: []*datapb.FieldBinlog{
{Binlogs: []*datapb.Binlog{{MemorySize: 100, LogSize: 80}}},
},
},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{
MmapVectorField: true,
})
require.NoError(t, err)
assert.Equal(t, int64(8100), usage.MemoryBytes)
assert.Equal(t, int64(512000), usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceDoesNotRequireRows(t *testing.T) {
schema := testSchema()
segments := []*datapb.SegmentInfo{
{
ID: 1,
Binlogs: []*datapb.FieldBinlog{
{FieldID: 1, Binlogs: []*datapb.Binlog{{MemorySize: 100, LogSize: 80}}},
},
},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{})
require.NoError(t, err)
assert.Equal(t, int64(100), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceSkipsRawVectorWhenIndexHasRawData(t *testing.T) {
schema := testSchema()
segments := []*datapb.SegmentInfo{{ID: 1, NumOfRows: 1000, Binlogs: testRawBinlogs()}}
indexes := map[int64][]*querypb.FieldIndexInfo{
1: {testVectorIndex("HNSW", 10000)},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, indexes, EstimateOptions{})
require.NoError(t, err)
assert.Equal(t, int64(18000), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceKeepsRawVectorWhenPreferFieldData(t *testing.T) {
schema := testSchema()
segments := []*datapb.SegmentInfo{{ID: 1, NumOfRows: 1000, Binlogs: testRawBinlogs()}}
indexes := map[int64][]*querypb.FieldIndexInfo{
1: {testVectorIndex("HNSW", 10000)},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, indexes, EstimateOptions{
PreferFieldDataWhenIndexHasRawData: true,
})
require.NoError(t, err)
assert.Equal(t, int64(530000), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceEstimatesHNSWIndexWhenSizeMissing(t *testing.T) {
schema := testSchema()
segments := []*datapb.SegmentInfo{{ID: 1, NumOfRows: 1000, Binlogs: testRawBinlogs()}}
indexes := map[int64][]*querypb.FieldIndexInfo{
1: {{
FieldID: 100,
NumRows: 1000,
IndexFilePaths: []string{
"/path/that/must/not/exist/index_file",
},
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "HNSW"},
{Key: common.MetricTypeKey, Value: "L2"},
{Key: common.DimKey, Value: "128"},
{Key: "M", Value: "8"},
},
}},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, indexes, EstimateOptions{})
require.NoError(t, err)
assert.Equal(t, int64(8000), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceUsesCIndexEstimate(t *testing.T) {
schema := testSchema()
segments := []*datapb.SegmentInfo{{ID: 1, NumOfRows: 1000, Binlogs: testRawBinlogs()}}
indexes := map[int64][]*querypb.FieldIndexInfo{
1: {testVectorIndex("HNSW", 10000)},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, indexes, EstimateOptions{})
require.NoError(t, err)
assert.Equal(t, int64(18000), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceReturnsErrorWhenIndexTypeMissing(t *testing.T) {
schema := testSchema()
segments := []*datapb.SegmentInfo{{ID: 1, NumOfRows: 1000, Binlogs: testRawBinlogs()}}
indexes := map[int64][]*querypb.FieldIndexInfo{
1: {{
FieldID: 100,
IndexSize: 10000,
IndexFilePaths: []string{"/path/that/must/not/exist/index_file"},
}},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, indexes, EstimateOptions{})
assert.ErrorIs(t, err, merr.ErrServiceInternal)
assert.Zero(t, usage)
}
func TestEstimateSegmentsLoadResourceKeepsFlatRawFieldWhenIndexSizeMissing(t *testing.T) {
schema := testSchema()
segments := []*datapb.SegmentInfo{{ID: 1, NumOfRows: 1000, Binlogs: testRawBinlogs()}}
indexes := map[int64][]*querypb.FieldIndexInfo{
1: {{
FieldID: 100,
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "FLAT"},
},
}},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, indexes, EstimateOptions{
MmapVectorField: true,
})
require.NoError(t, err)
assert.Equal(t, int64(8000), usage.MemoryBytes)
assert.Equal(t, int64(512000), usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceIncludesJSONKeyStats(t *testing.T) {
schema := &schemapb.CollectionSchema{}
segments := []*datapb.SegmentInfo{
{
ID: 1,
NumOfRows: 1000,
JsonKeyStats: map[int64]*datapb.JsonKeyStats{
101: {MemorySize: 1000, LogSize: 800},
},
},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{
JSONKeyStatsExpansionFactor: 1.5,
})
require.NoError(t, err)
assert.Equal(t, int64(1500), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
usage, err = EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{
MmapJSONStats: true,
JSONKeyStatsExpansionFactor: 1.5,
})
require.NoError(t, err)
assert.Zero(t, usage.MemoryBytes)
assert.Equal(t, int64(1500), usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceIncludesTextStats(t *testing.T) {
schema := &schemapb.CollectionSchema{}
segments := []*datapb.SegmentInfo{
{
ID: 1,
NumOfRows: 1000,
TextStatsLogs: map[int64]*datapb.TextIndexStats{
101: {MemorySize: 2000, LogSize: 1500},
},
},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{
TextIndexExpansionFactor: 1.25,
})
require.NoError(t, err)
// 1000 rows also reserve a 128-byte word-aligned Tantivy validity bitmap.
assert.Equal(t, int64(2628), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
usage, err = EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{
MmapScalarField: true,
TextIndexExpansionFactor: 1.25,
})
require.NoError(t, err)
assert.Equal(t, int64(128), usage.MemoryBytes)
assert.Equal(t, int64(2500), usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceUsesStorageV3Descriptor(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{{FieldID: 100, Name: "value", DataType: schemapb.DataType_Int64}},
}
segments := []*datapb.SegmentInfo{{
ID: 1,
NumOfRows: 10,
StorageVersion: storage.StorageV3,
ManifestPath: "files/manifest",
Stats: &datapb.Statistics{
LoadResource: &datapb.LoadResourceStatistics{
ColumnGroups: []*datapb.ColumnGroupStatistics{{
GroupId: 100,
FieldIds: []int64{100},
MemorySize: 4096,
}},
},
},
}}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{})
require.NoError(t, err)
assert.Equal(t, int64(4096), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceAppliesTieredRatioToStats(t *testing.T) {
schema := &schemapb.CollectionSchema{}
segments := []*datapb.SegmentInfo{
{
ID: 1,
NumOfRows: 1000,
JsonKeyStats: map[int64]*datapb.JsonKeyStats{
101: {MemorySize: 101},
},
TextStatsLogs: map[int64]*datapb.TextIndexStats{
102: {MemorySize: 201},
},
},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{
MmapScalarField: true,
TieredEvictionEnabled: true,
TieredMemoryCacheRatio: 0.3,
TieredDiskCacheRatio: 0.5,
})
require.NoError(t, err)
// The memory ratio applies to JSON stats plus the text validity bitmap.
assert.Equal(t, int64(69), usage.MemoryBytes)
assert.Equal(t, int64(101), usage.DiskBytes)
}
func TestEstimateSegmentsLoadResourceAllowsZeroTieredRatio(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{
FieldID: 101,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "4"},
},
},
},
}
segments := []*datapb.SegmentInfo{
{
ID: 1,
NumOfRows: 10,
Binlogs: []*datapb.FieldBinlog{
{
FieldID: 100,
Binlogs: []*datapb.Binlog{
{MemorySize: 100, LogSize: 80},
},
},
{
FieldID: 101,
Binlogs: []*datapb.Binlog{
{MemorySize: 200, LogSize: 160},
},
},
},
Statslogs: []*datapb.FieldBinlog{
{Binlogs: []*datapb.Binlog{{MemorySize: 30, LogSize: 20}}},
},
},
}
usage, err := EstimateSegmentsLoadResource(context.Background(), schema, segments, nil, EstimateOptions{
MmapVectorField: true,
TieredEvictionEnabled: true,
TieredMemoryCacheRatio: 0,
TieredDiskCacheRatio: 0,
})
require.NoError(t, err)
assert.Equal(t, int64(30), usage.MemoryBytes)
assert.Zero(t, usage.DiskBytes)
}
func TestFilterSegmentsByLoadFieldsKeepsChildFieldBinlogs(t *testing.T) {
segments := []*datapb.SegmentInfo{
{
ID: 1,
Binlogs: []*datapb.FieldBinlog{
{FieldID: 10, ChildFields: []int64{100, 101}},
{FieldID: 11, ChildFields: []int64{102}},
{FieldID: 200},
{FieldID: 201},
},
Stats: &datapb.Statistics{
LoadResource: &datapb.LoadResourceStatistics{
ColumnGroups: []*datapb.ColumnGroupStatistics{
{GroupId: 10, FieldIds: []int64{100, 101}},
{GroupId: 11, FieldIds: []int64{102}},
{GroupId: 200},
{GroupId: 201},
},
},
},
},
}
filtered := FilterSegmentsByLoadFields(segments, []int64{100, 200})
require.Len(t, filtered, 1)
require.Len(t, filtered[0].GetBinlogs(), 2)
assert.Equal(t, int64(10), filtered[0].GetBinlogs()[0].GetFieldID())
assert.Equal(t, int64(200), filtered[0].GetBinlogs()[1].GetFieldID())
require.Len(t, filtered[0].GetStats().GetLoadResource().GetColumnGroups(), 2)
assert.Equal(t, int64(10), filtered[0].GetStats().GetLoadResource().GetColumnGroups()[0].GetGroupId())
assert.Equal(t, int64(200), filtered[0].GetStats().GetLoadResource().GetColumnGroups()[1].GetGroupId())
}
func TestFilterLoadFieldsKeepsSystemFields(t *testing.T) {
segments := []*datapb.SegmentInfo{
{
ID: 1,
Binlogs: []*datapb.FieldBinlog{
{FieldID: common.RowIDField},
{FieldID: common.TimeStampField},
{FieldID: 100},
{FieldID: 101},
},
Statslogs: []*datapb.FieldBinlog{
{FieldID: common.RowIDField},
{FieldID: 100},
{FieldID: 101},
},
},
}
filteredSegments := FilterSegmentsByLoadFields(segments, []int64{100})
require.Len(t, filteredSegments, 1)
assert.ElementsMatch(t, []int64{common.RowIDField, common.TimeStampField, 100}, lo.Map(filteredSegments[0].GetBinlogs(), func(fieldBinlog *datapb.FieldBinlog, _ int) int64 {
return fieldBinlog.GetFieldID()
}))
assert.ElementsMatch(t, []int64{common.RowIDField, 100}, lo.Map(filteredSegments[0].GetStatslogs(), func(fieldBinlog *datapb.FieldBinlog, _ int) int64 {
return fieldBinlog.GetFieldID()
}))
}
func TestEstimateSegmentsLoadResourceForLoadConfigKeepsFullSchemaForChildFields(t *testing.T) {
schema := &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "loaded", DataType: schemapb.DataType_Int64},
{FieldID: 101, Name: "sibling", DataType: schemapb.DataType_Int64},
},
}
segments := []*datapb.SegmentInfo{
{
ID: 1,
NumOfRows: 100,
Binlogs: []*datapb.FieldBinlog{
{
FieldID: 10,
ChildFields: []int64{100, 101},
Binlogs: []*datapb.Binlog{
{MemorySize: 100, LogSize: 80},
},
},
},
},
}
usage, err := EstimateSegmentsLoadResourceForLoadConfig(context.Background(), schema, segments, nil, []int64{100}, EstimateOptions{
MmapScalarField: true,
})
require.NoError(t, err)
assert.Zero(t, usage.MemoryBytes)
assert.Equal(t, int64(100), usage.DiskBytes)
}
func testRawBinlogs() []*datapb.FieldBinlog {
return []*datapb.FieldBinlog{
{
FieldID: 1,
Binlogs: []*datapb.Binlog{
{MemorySize: 8000, LogSize: 8000},
},
},
{
FieldID: 100,
Binlogs: []*datapb.Binlog{
{MemorySize: 512000, LogSize: 512000},
},
},
}
}
func testVectorIndex(indexType string, indexSize int64) *querypb.FieldIndexInfo {
return &querypb.FieldIndexInfo{
FieldID: 100,
IndexSize: indexSize,
NumRows: 1000,
IndexFilePaths: []string{
"/path/that/must/not/exist/index_file",
},
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: indexType},
{Key: common.MetricTypeKey, Value: "L2"},
{Key: common.DimKey, Value: "128"},
},
}
}
func testSchema() *schemapb.CollectionSchema {
return &schemapb.CollectionSchema{
Fields: []*schemapb.FieldSchema{
{FieldID: 1, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{
FieldID: 100,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "128"},
},
},
},
}
}
internal/querycoordv2/autoscale_precheck.go
+636
View File
@@ -0,0 +1,636 @@
// 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 querycoordv2
import (
"context"
"strconv"
"time"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"golang.org/x/sync/errgroup"
"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/internal/querycoordv2/autoscale"
"github.com/milvus-io/milvus/internal/querycoordv2/job"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/internal/querycoordv2/task"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"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/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/metricsinfo"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/retry"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
const (
loadResourcePrecheckSuggestKey = "suggested_expand_percent"
maxSegmentNumPerGetIndexInfoRequest = 1024
maxConcurrentQueryNodeMetrics = 16
queryNodeMetricsTimeout = 10 * time.Second
bytesPerMB = 1024 * 1024
bytesPerGB = 1024 * 1024 * 1024
)
type autoscaleResourceUsage struct {
memoryBytes int64
diskBytes int64
}
type queryNodeResourceStatus struct {
available autoscaleResourceUsage
capacity autoscaleResourceUsage // capacity before memory and disk admission thresholds
}
func (usage autoscaleResourceUsage) isZero() bool {
return usage.memoryBytes == 0 && usage.diskBytes == 0
}
func (usage autoscaleResourceUsage) multiply(factor int64) autoscaleResourceUsage {
return autoscaleResourceUsage{
memoryBytes: usage.memoryBytes * factor,
diskBytes: usage.diskBytes * factor,
}
}
func (usage *autoscaleResourceUsage) add(other autoscaleResourceUsage) {
usage.memoryBytes += other.memoryBytes
usage.diskBytes += other.diskBytes
}
func (usage autoscaleResourceUsage) shortageFrom(current autoscaleResourceUsage) autoscaleResourceUsage {
return autoscaleResourceUsage{
memoryBytes: max(usage.memoryBytes-current.memoryBytes, 0),
diskBytes: max(usage.diskBytes-current.diskBytes, 0),
}
}
type loadResourcePrecheckError struct {
message string
suggestedExpandPercent int64
}
func (e *loadResourcePrecheckError) Error() string {
return e.message
}
func (e *loadResourcePrecheckError) Unwrap() error {
return merr.ErrServiceResourceInsufficient
}
func statusWithLoadResourcePrecheckSuggestion(err error) *commonpb.Status {
status := merr.Status(err)
var precheckErr *loadResourcePrecheckError
if !errors.As(err, &precheckErr) {
return status
}
status.Retriable = false
if status.ExtraInfo == nil {
status.ExtraInfo = make(map[string]string)
}
status.ExtraInfo[loadResourcePrecheckSuggestKey] = strconv.FormatInt(precheckErr.suggestedExpandPercent, 10)
return status
}
func checkAutoscaleLoadResourceByRGWithCapacity(requiredByRG, availableByRG, capacityByRG map[string]autoscaleResourceUsage, globalCapacity, globalLimit autoscaleResourceUsage, autoscaleEnabled bool) error {
totalShortage := autoscaleResourceUsage{}
suggestedExpandPercent := int64(0)
for rgName, required := range requiredByRG {
available := availableByRG[rgName]
if required.memoryBytes <= available.memoryBytes && required.diskBytes <= available.diskBytes {
continue
}
shortage := required.shortageFrom(available)
totalShortage.add(shortage)
capacity := capacityByRG[rgName]
suggestedExpandPercent = max(
suggestedExpandPercent,
max(
calculateExpandPercentFromShortage(shortage.memoryBytes, capacity.memoryBytes),
calculateExpandPercentFromShortage(shortage.diskBytes, capacity.diskBytes),
),
)
}
if totalShortage.isZero() {
return nil
}
if autoscaleEnabled && canExpandWithinGlobalLimit(totalShortage, globalCapacity, globalLimit) {
return nil
}
return newLoadResourcePrecheckError(suggestedExpandPercent)
}
func canExpandWithinGlobalLimit(shortage, globalCapacity, globalLimit autoscaleResourceUsage) bool {
remainingCapacity := autoscaleResourceUsage{
memoryBytes: max(globalLimit.memoryBytes-globalCapacity.memoryBytes, 0),
diskBytes: max(globalLimit.diskBytes-globalCapacity.diskBytes, 0),
}
availableAfterScale := usableResourceCapacity(remainingCapacity)
return shortage.memoryBytes <= availableAfterScale.memoryBytes &&
shortage.diskBytes <= availableAfterScale.diskBytes
}
func newLoadResourcePrecheckError(suggestedExpandPercent int64) error {
return &loadResourcePrecheckError{
message: "Insufficient cluster resources. Please scale out the cluster or enable autoscaling. Suggested scale-out ratio: " + strconv.FormatInt(suggestedExpandPercent, 10) + "%",
suggestedExpandPercent: suggestedExpandPercent,
}
}
func newQueryNodeResourceMetricsUnavailableError(scope string) error {
return merr.WrapErrServiceUnavailableMsg("query node resource metrics unavailable, scope=%s", scope)
}
func calculateExpandPercentFromShortage(shortage, capacity int64) int64 {
if shortage <= 0 {
return 0
}
if capacity <= 0 {
return 100
}
return (shortage*100 + capacity - 1) / capacity
}
func (s *Server) checkLoadResource(ctx context.Context, req *job.AlterLoadConfigRequest) (map[string]autoscaleResourceUsage, error) {
if !paramtable.Get().QueryCoordCfg.AutoscalePrecheckEnabled.GetAsBool() {
return nil, nil
}
requiredByRG, err := s.evaluateLoadResource(ctx, req)
if err == nil {
return requiredByRG, nil
}
var precheckErr *loadResourcePrecheckError
if errors.As(err, &precheckErr) {
return nil, err
}
mlog.Warn(ctx, "skip load resource precheck due to internal error",
mlog.Int64("collectionID", req.CollectionInfo.GetCollectionID()),
mlog.Err(err),
)
return requiredByRG, nil
}
func (s *Server) evaluateLoadResource(ctx context.Context, req *job.AlterLoadConfigRequest) (map[string]autoscaleResourceUsage, error) {
requiredByRG, err := s.requiredLoadResourceByRG(ctx, req)
if err != nil {
return nil, err
}
if len(requiredByRG) == 0 {
return nil, nil
}
// Physical usage includes evictable caches and cannot establish whether a
// tiered node can admit this load. Preserve demand for the broadcast counter.
if paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool() {
return requiredByRG, nil
}
autoscaleEnabled := paramtable.Get().QueryCoordCfg.AutoscaleEnabled.GetAsBool()
currentByRG, capacityByRG, globalCapacity, err := s.currentQueryNodeResourceStatus(ctx, requiredByRG, autoscaleEnabled)
if err != nil {
return requiredByRG, err
}
globalLimit := autoscaleGlobalLimitFromConfig(
paramtable.Get().QueryCoordCfg.AutoscaleMaxMemoryLimit.GetAsInt64(),
paramtable.Get().QueryCoordCfg.AutoscaleMaxDiskLimit.GetAsInt64(),
)
err = checkAutoscaleLoadResourceByRGWithCapacity(requiredByRG, currentByRG, capacityByRG, globalCapacity, globalLimit, autoscaleEnabled)
mlog.Info(ctx, "autoscale load resource precheck",
mlog.Int64("collectionID", req.CollectionInfo.GetCollectionID()),
mlog.Bool("autoscaleEnabled", autoscaleEnabled),
mlog.Bool("passed", err == nil),
mlog.Any("requiredByRGMB", resourceUsageByRGToMB(requiredByRG)),
mlog.Any("availableByRGMB", resourceUsageByRGToMB(currentByRG)),
mlog.Any("shortageByRGMB", shortageByRGToMB(requiredByRG, currentByRG)),
mlog.Any("capacityByRGMB", resourceUsageByRGToMB(capacityByRG)),
mlog.Float64("globalCapacityMemoryMB", bytesToMB(globalCapacity.memoryBytes)),
mlog.Float64("globalCapacityDiskMB", bytesToMB(globalCapacity.diskBytes)),
mlog.Float64("globalLimitMemoryMB", bytesToMB(globalLimit.memoryBytes)),
mlog.Float64("globalLimitDiskMB", bytesToMB(globalLimit.diskBytes)),
mlog.Err(err),
)
if err != nil {
return nil, err
}
return requiredByRG, nil
}
func recordLoadResourceDemand(requiredByRG map[string]autoscaleResourceUsage) {
for rgName, usage := range requiredByRG {
metrics.QueryCoordLoadDemandMemoryBytes.WithLabelValues(rgName).Add(float64(usage.memoryBytes))
metrics.QueryCoordLoadDemandDiskBytes.WithLabelValues(rgName).Add(float64(usage.diskBytes))
}
}
func resourceUsageByRGToMB(usages map[string]autoscaleResourceUsage) map[string]map[string]float64 {
result := make(map[string]map[string]float64, len(usages))
for rgName, usage := range usages {
result[rgName] = map[string]float64{
"memory": bytesToMB(usage.memoryBytes),
"disk": bytesToMB(usage.diskBytes),
}
}
return result
}
func shortageByRGToMB(requiredByRG, currentByRG map[string]autoscaleResourceUsage) map[string]map[string]float64 {
shortageByRG := make(map[string]autoscaleResourceUsage, len(requiredByRG))
for rgName, required := range requiredByRG {
shortageByRG[rgName] = required.shortageFrom(currentByRG[rgName])
}
return resourceUsageByRGToMB(shortageByRG)
}
func bytesToMB(bytes int64) float64 {
return float64(bytes) / bytesPerMB
}
func autoscaleGlobalLimitFromConfig(maxMemoryLimitGB, maxDiskLimitGB int64) autoscaleResourceUsage {
return autoscaleResourceUsage{
memoryBytes: maxMemoryLimitGB * bytesPerGB,
diskBytes: maxDiskLimitGB * bytesPerGB,
}
}
func (s *Server) requiredLoadResourceByRG(ctx context.Context, req *job.AlterLoadConfigRequest) (map[string]autoscaleResourceUsage, error) {
collectionID := req.CollectionInfo.GetCollectionID()
expectedPartitionIDs := req.Expected.ExpectedPartitionIDs
if len(expectedPartitionIDs) == 0 {
return nil, nil
}
_, segments, err := s.broker.GetRecoveryInfoV2(ctx, collectionID, expectedPartitionIDs...)
if err != nil {
return nil, err
}
if len(segments) == 0 {
return nil, nil
}
segmentIDs := lo.Map(segments, func(segment *datapb.SegmentInfo, _ int) int64 {
return segment.GetID()
})
// Recovery info only carries compact segment metadata, while resource
// estimation needs the complete binlog and statistics payload.
segments, err = s.broker.GetSegmentInfo(ctx, segmentIDs...)
if err != nil {
return nil, err
}
indexes, err := getIndexInfoBySegments(ctx, s.broker, collectionID, segmentIDs)
if err != nil {
return nil, err
}
expectedUsage, err := estimateLoadResourceForLoadConfig(
ctx,
req.CollectionInfo.GetSchema(),
req.CollectionInfo.GetProperties(),
segments,
indexes,
req.Expected.ExpectedLoadFields,
)
if err != nil {
return nil, err
}
if expectedUsage.isZero() {
return nil, nil
}
currentUsage := autoscaleResourceUsage{}
currentReplicasNumber := map[string]int{}
if req.Current.Collection != nil {
currentPartitionSet := typeutil.NewSet(req.Current.GetPartitionIDs()...)
currentSegments := filterSegmentsByPartitionIDs(segments, currentPartitionSet)
if len(currentSegments) > 0 {
currentUsage, err = estimateLoadResourceForLoadConfig(
ctx,
currentSchema(req),
req.CollectionInfo.GetProperties(),
currentSegments,
indexes,
req.Current.GetLoadFields(),
)
if err != nil {
return nil, err
}
}
currentReplicasNumber = req.Current.GetReplicaNumber()
}
return buildRequiredLoadResourceByRGForLoadConfig(expectedUsage, currentUsage, currentReplicasNumber, req.Expected.ExpectedReplicaNumber), nil
}
func currentSchema(req *job.AlterLoadConfigRequest) *schemapb.CollectionSchema {
if req.Current.Collection != nil && req.Current.Collection.Schema != nil {
return req.Current.Collection.Schema
}
return req.CollectionInfo.GetSchema()
}
func estimateLoadResourceForLoadConfig(ctx context.Context, schema *schemapb.CollectionSchema, collectionProperties []*commonpb.KeyValuePair, segments []*datapb.SegmentInfo, indexes map[int64][]*querypb.FieldIndexInfo, loadFields []int64) (autoscaleResourceUsage, error) {
if schema != nil {
schema = task.ApplyCollectionMmapSetting(typeutil.Clone(schema), collectionProperties)
}
usage, err := autoscale.EstimateSegmentsLoadResourceForLoadConfig(ctx, schema, segments, indexes, loadFields, autoscale.DefaultEstimateOptions())
if err != nil {
return autoscaleResourceUsage{}, err
}
return autoscaleResourceUsage{
memoryBytes: usage.MemoryBytes,
diskBytes: usage.DiskBytes,
}, nil
}
func getIndexInfoBySegments(ctx context.Context, broker meta.Broker, collectionID int64, segmentIDs []int64) (map[int64][]*querypb.FieldIndexInfo, error) {
if len(segmentIDs) == 0 {
return nil, nil
}
indexes := make(map[int64][]*querypb.FieldIndexInfo)
for _, batch := range lo.Chunk(segmentIDs, maxSegmentNumPerGetIndexInfoRequest) {
batchIndexes, err := broker.GetIndexInfo(ctx, collectionID, batch...)
if err != nil {
if errors.Is(err, merr.ErrIndexNotFound) {
continue
}
return nil, err
}
for segmentID, infos := range batchIndexes {
indexes[segmentID] = infos
}
}
if len(indexes) == 0 {
return nil, nil
}
return indexes, nil
}
func filterSegmentsByPartitionIDs(segments []*datapb.SegmentInfo, partitionIDs typeutil.Set[int64]) []*datapb.SegmentInfo {
if len(partitionIDs) == 0 {
return nil
}
return lo.Filter(segments, func(segment *datapb.SegmentInfo, _ int) bool {
return partitionIDs.Contain(segment.GetPartitionID())
})
}
func buildRequiredLoadResourceByRG(required autoscaleResourceUsage, expectedReplicasNumber map[string]int) map[string]autoscaleResourceUsage {
requiredByRG := make(map[string]autoscaleResourceUsage)
for rgName, num := range expectedReplicasNumber {
if num <= 0 {
continue
}
requiredByRG[rgName] = required.multiply(int64(num))
}
if len(requiredByRG) == 0 {
requiredByRG[meta.DefaultResourceGroupName] = required
}
return requiredByRG
}
func buildRequiredLoadResourceByRGForLoadConfig(expectedUsage, currentUsage autoscaleResourceUsage, currentReplicasNumber, expectedReplicasNumber map[string]int) map[string]autoscaleResourceUsage {
if len(expectedReplicasNumber) == 0 {
return buildRequiredLoadResourceByRG(expectedUsage, expectedReplicasNumber)
}
loadDeltaForKeptReplica := expectedUsage.shortageFrom(currentUsage)
requiredByRG := make(map[string]autoscaleResourceUsage)
for rgName, expectedNum := range expectedReplicasNumber {
if expectedNum <= 0 {
continue
}
currentNum := currentReplicasNumber[rgName]
required := autoscaleResourceUsage{}
keptNum := min(currentNum, expectedNum)
if keptNum > 0 {
required.add(loadDeltaForKeptReplica.multiply(int64(keptNum)))
}
incomingNum := expectedNum - currentNum
if incomingNum > 0 {
required.add(expectedUsage.multiply(int64(incomingNum)))
}
if !required.isZero() {
requiredByRG[rgName] = required
}
}
return requiredByRG
}
func (s *Server) currentQueryNodeResourceStatus(ctx context.Context, requiredByRG map[string]autoscaleResourceUsage, includeGlobal bool) (map[string]autoscaleResourceUsage, map[string]autoscaleResourceUsage, autoscaleResourceUsage, error) {
nodesByRG := make(map[string][]*session.NodeInfo, len(requiredByRG))
nodesByID := make(map[int64]*session.NodeInfo)
for rgName := range requiredByRG {
nodes, err := s.queryNodesInResourceGroup(ctx, rgName)
if err != nil {
return nil, nil, autoscaleResourceUsage{}, err
}
nodesByRG[rgName] = nodes
for _, node := range nodes {
nodesByID[node.ID()] = node
}
}
globalNodes := make([]*session.NodeInfo, 0)
if includeGlobal {
globalNodes = s.onlineQueryNodes()
for _, node := range globalNodes {
nodesByID[node.ID()] = node
}
}
snapshot, err := s.queryNodeResourceSnapshot(ctx, lo.Values(nodesByID))
if err != nil {
return nil, nil, autoscaleResourceUsage{}, err
}
currentByRG := make(map[string]autoscaleResourceUsage, len(requiredByRG))
capacityByRG := make(map[string]autoscaleResourceUsage, len(requiredByRG))
for rgName, nodes := range nodesByRG {
available, capacity, ok := queryNodeResourceStatusFromSnapshot(nodes, snapshot)
if !ok {
return nil, nil, autoscaleResourceUsage{}, newQueryNodeResourceMetricsUnavailableError("resource-group/" + rgName)
}
currentByRG[rgName] = available
capacityByRG[rgName] = usableResourceCapacity(capacity)
}
globalCapacity := autoscaleResourceUsage{}
if includeGlobal {
_, capacity, ok := queryNodeResourceStatusFromSnapshot(globalNodes, snapshot)
if !ok {
return nil, nil, autoscaleResourceUsage{}, newQueryNodeResourceMetricsUnavailableError("global")
}
globalCapacity = capacity
}
return currentByRG, capacityByRG, globalCapacity, nil
}
func (s *Server) queryNodesInResourceGroup(ctx context.Context, rgName string) ([]*session.NodeInfo, error) {
nodeIDs, err := s.meta.GetNodes(ctx, rgName)
if err != nil {
return nil, err
}
nodes := make([]*session.NodeInfo, 0, len(nodeIDs))
for _, nodeID := range nodeIDs {
node := s.nodeMgr.Get(nodeID)
if node == nil || node.IsStoppingState() {
continue
}
nodes = append(nodes, node)
}
return nodes, nil
}
func (s *Server) onlineQueryNodes() []*session.NodeInfo {
if s.nodeMgr == nil {
return nil
}
return lo.Filter(s.nodeMgr.GetAll(), func(node *session.NodeInfo, _ int) bool {
return node != nil && !node.IsStoppingState() && !node.IsEmbeddedQueryNodeInStreamingNode()
})
}
func (s *Server) queryNodeResourceSnapshot(ctx context.Context, nodes []*session.NodeInfo) (map[int64]queryNodeResourceStatus, error) {
snapshot := make(map[int64]queryNodeResourceStatus, len(nodes))
if len(nodes) == 0 {
return snapshot, nil
}
if s.cluster == nil {
diskCapacityPerNode := paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsInt64()
for _, node := range nodes {
capacity := autoscaleResourceUsage{diskBytes: diskCapacityPerNode}
if node.MemCapacity() > 0 {
capacity.memoryBytes = int64(node.MemCapacity() * 1024 * 1024)
}
snapshot[node.ID()] = queryNodeResourceStatus{available: usableResourceCapacity(capacity), capacity: capacity}
}
return snapshot, nil
}
req, err := metricsinfo.ConstructRequestByMetricType(metricsinfo.SystemInfoMetrics)
if err != nil {
return nil, merr.WrapErrServiceInternalErr(err, "failed to construct query node resource metrics request")
}
statuses := make([]queryNodeResourceStatus, len(nodes))
metricsCtx, cancel := context.WithTimeout(ctx, queryNodeMetricsTimeout)
defer cancel()
metricsCtx = retry.WithMaxAttemptsContext(metricsCtx, 1)
group, metricsCtx := errgroup.WithContext(metricsCtx)
group.SetLimit(min(maxConcurrentQueryNodeMetrics, len(nodes)))
for index, node := range nodes {
index, node := index, node
group.Go(func() error {
resp, err := s.cluster.GetMetrics(metricsCtx, node.ID(), req)
if err := merr.CheckRPCCall(resp, err); err != nil {
mlog.Warn(ctx, "failed to get query node resource metrics",
mlog.Int64("nodeID", node.ID()),
mlog.Err(err),
)
return merr.WrapErrServiceUnavailableMsg("query node resource metrics unavailable, nodeID=%d", node.ID())
}
nodeInfo := &metricsinfo.QueryNodeInfos{}
if err := metricsinfo.UnmarshalComponentInfos(resp.GetResponse(), nodeInfo); err != nil {
mlog.Warn(ctx, "failed to parse query node resource metrics",
mlog.Int64("nodeID", node.ID()),
mlog.Err(err),
)
return merr.WrapErrServiceInternalMsg("failed to parse query node resource metrics, nodeID=%d", node.ID())
}
hardware := nodeInfo.HardwareInfos
if hardware.Memory == 0 {
mlog.Warn(ctx, "query node resource metrics missing memory capacity",
mlog.Int64("nodeID", node.ID()),
)
return merr.WrapErrServiceInternalMsg("query node resource metrics missing memory capacity, nodeID=%d", node.ID())
}
availableMemoryBytes, capacityMemoryBytes := memoryResourceFromMetrics(hardware.Memory, hardware.MemoryUsage)
localStorage := nodeInfo.LocalStorage
if localStorage == nil || localStorage.CapacityBytes <= 0 || localStorage.UsedBytes < 0 {
return merr.WrapErrServiceInternalMsg("query node resource metrics missing local storage usage, nodeID=%d", node.ID())
}
usableLocalDiskBytes := usableResourceCapacity(autoscaleResourceUsage{diskBytes: localStorage.CapacityBytes}).diskBytes
statuses[index] = queryNodeResourceStatus{
available: autoscaleResourceUsage{
memoryBytes: availableMemoryBytes,
diskBytes: max(usableLocalDiskBytes-localStorage.UsedBytes, 0),
},
capacity: autoscaleResourceUsage{memoryBytes: capacityMemoryBytes, diskBytes: localStorage.CapacityBytes},
}
return nil
})
}
if err := group.Wait(); err != nil {
return nil, err
}
for index, node := range nodes {
snapshot[node.ID()] = statuses[index]
}
return snapshot, nil
}
func queryNodeResourceStatusFromSnapshot(nodes []*session.NodeInfo, snapshot map[int64]queryNodeResourceStatus) (autoscaleResourceUsage, autoscaleResourceUsage, bool) {
if len(nodes) == 0 {
return autoscaleResourceUsage{}, autoscaleResourceUsage{}, false
}
available := autoscaleResourceUsage{}
capacity := autoscaleResourceUsage{}
for _, node := range nodes {
status, ok := snapshot[node.ID()]
if !ok {
return autoscaleResourceUsage{}, autoscaleResourceUsage{}, false
}
available.add(status.available)
capacity.add(status.capacity)
}
return available, capacity, true
}
func memoryResourceFromMetrics(totalMemory, usedMemory uint64) (int64, int64) {
memoryCapacity := uint64ToInt64(totalMemory)
memoryLimit := usableResourceCapacity(autoscaleResourceUsage{memoryBytes: memoryCapacity}).memoryBytes
return memoryLimit - uint64ToInt64(usedMemory), memoryCapacity
}
func usableResourceCapacity(capacity autoscaleResourceUsage) autoscaleResourceUsage {
return autoscaleResourceUsage{
memoryBytes: int64(float64(capacity.memoryBytes) * paramtable.Get().QuotaConfig.QueryNodeMemoryHighWaterLevel.GetAsFloat()),
diskBytes: int64(float64(capacity.diskBytes) * paramtable.Get().QueryNodeCfg.MaxDiskUsagePercentage.GetAsFloat()),
}
}
func uint64ToInt64(value uint64) int64 {
const maxInt64 = uint64(1<<63 - 1)
if value > maxInt64 {
return int64(maxInt64)
}
return int64(value)
}
internal/querycoordv2/autoscale_precheck_test.go
+1271
View File
File diff suppressed because it is too large Load Diff
internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go
+9 -1
View File
@@ -126,8 +126,16 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadCollection(ctx conte
mlog.Int64("collectionID", req.GetCollectionID()))
return nil
}
requiredByRG, err := s.checkLoadResource(ctx, alterLoadConfigReq)
if err != nil {
return err
}
_, err = broadcaster.Broadcast(ctx, msg)
return err
if err != nil {
return err
}
recordLoadResourceDemand(requiredByRG)
return nil
}
// getLoadReplicaConfigForRequest resolves the replica configuration a load
internal/querycoordv2/ddl_callbacks_alter_load_info_load_partitions.go
+10 -2
View File
@@ -90,7 +90,7 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadPartitions(ctx conte
alterLoadConfigReq := &job.AlterLoadConfigRequest{
Meta: s.meta,
CollectionInfo: coll,
Current: s.getCurrentLoadConfig(ctx, req.GetCollectionID()),
Current: currentLoadConfig,
Expected: job.ExpectedLoadConfig{
ExpectedPartitionIDs: partitionIDsSet.Collect(),
ExpectedReplicaNumber: expectedReplicasNumber,
@@ -115,6 +115,14 @@ func (s *Server) broadcastAlterLoadConfigCollectionV2ForLoadPartitions(ctx conte
mlog.Int64s("partitionIDs", req.GetPartitionIDs()))
return nil
}
requiredByRG, err := s.checkLoadResource(ctx, alterLoadConfigReq)
if err != nil {
return err
}
_, err = broadcaster.Broadcast(ctx, msg)
return err
if err != nil {
return err
}
recordLoadResourceDemand(requiredByRG)
return nil
}
internal/querycoordv2/ddl_callbacks_alter_load_info_test.go
+113
View File
@@ -22,15 +22,128 @@ import (
"github.com/bytedance/mockey"
"github.com/cockroachdb/errors"
"github.com/prometheus/client_golang/prometheus/testutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/milvuspb"
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_broadcaster"
"github.com/milvus-io/milvus/internal/querycoordv2/job"
"github.com/milvus-io/milvus/internal/querycoordv2/meta"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/proto/messagespb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
func TestLoadCallbacksRecordDemandAfterBroadcast(t *testing.T) {
enableAutoscalePrecheckForTest(t)
setAutoscaleCapacityRatiosForTest(t, "0.9", "0.5")
params := paramtable.Get()
require.NoError(t, params.Save(params.QueryCoordCfg.AutoscaleEnabled.Key, "false"))
t.Cleanup(func() { params.Reset(params.QueryCoordCfg.AutoscaleEnabled.Key) })
for _, loadPartitions := range []bool{false, true} {
operation := "LoadCollection"
if loadPartitions {
operation = "LoadPartitions"
}
for _, scenario := range []string{"precheck rejection", "broadcast failure", "success", "tiered success", "metrics failure", "metrics and broadcast failure"} {
mockey.PatchConvey(operation+"/"+scenario, t, func() {
tiered := scenario == "tiered success"
if tiered {
require.NoError(t, params.Save(params.QueryNodeCfg.TieredEvictionEnabled.Key, "true"))
defer params.Save(params.QueryNodeCfg.TieredEvictionEnabled.Key, "false")
}
ctx := context.Background()
cluster := session.NewMockCluster(t)
server := newAutoscalePrecheckTestServer(t, cluster)
broker := meta.NewMockBroker(t)
server.broker = broker
broker.EXPECT().DescribeCollection(mock.Anything, autoscalePrecheckTestCollectionID).
Return(newAutoscalePrecheckTestRequest().CollectionInfo, nil).Once()
if !loadPartitions {
broker.EXPECT().GetPartitions(mock.Anything, autoscalePrecheckTestCollectionID).
Return([]int64{autoscalePrecheckTestPartitionID}, nil).Once()
}
totalMemory := uint64(8192)
if scenario == "precheck rejection" {
totalMemory = 1024
}
if scenario == "metrics failure" || scenario == "metrics and broadcast failure" {
cluster.EXPECT().GetMetrics(mock.Anything, int64(1), mock.Anything).
Return(nil, errors.New("metrics unavailable")).Once()
} else if !tiered {
cluster.EXPECT().GetMetrics(mock.Anything, int64(1), mock.Anything).
Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: newAutoscalePrecheckMetrics(t, totalMemory, 0, 1)}, nil).Once()
}
mockey.Mock((*Server).requiredLoadResourceByRG).Return(map[string]autoscaleResourceUsage{
meta.DefaultResourceGroupName: {memoryBytes: 1024, diskBytes: 2048},
}, nil).Build()
msg := message.NewAlterLoadConfigMessageBuilderV2().
WithHeader(&messagespb.AlterLoadConfigMessageHeader{CollectionId: autoscalePrecheckTestCollectionID}).
WithBody(&messagespb.AlterLoadConfigMessageBody{}).
WithBroadcast([]string{"_ctrl_channel"}).MustBuildBroadcast()
mockey.Mock(job.GenerateAlterLoadConfigMessage).Return(msg, nil).Build()
bapi := mock_broadcaster.NewMockBroadcastAPI(t)
bapi.EXPECT().Close().Return().Once()
mockey.Mock((*Server).startBroadcastWithCollectionIDLock).Return(bapi, nil).Build()
memoryCounter := metrics.QueryCoordLoadDemandMemoryBytes.WithLabelValues(meta.DefaultResourceGroupName)
diskCounter := metrics.QueryCoordLoadDemandDiskBytes.WithLabelValues(meta.DefaultResourceGroupName)
memoryBefore := testutil.ToFloat64(memoryCounter)
diskBefore := testutil.ToFloat64(diskCounter)
broadcastErr := errors.New("broadcast failed")
if scenario != "precheck rejection" {
bapi.EXPECT().Broadcast(mock.Anything, msg).
RunAndReturn(func(context.Context, message.BroadcastMutableMessage) (*types.BroadcastAppendResult, error) {
assert.Equal(t, memoryBefore, testutil.ToFloat64(memoryCounter))
assert.Equal(t, diskBefore, testutil.ToFloat64(diskCounter))
if scenario == "broadcast failure" || scenario == "metrics and broadcast failure" {
return nil, broadcastErr
}
return &types.BroadcastAppendResult{}, nil
}).Once()
}
var err error
if loadPartitions {
err = server.broadcastAlterLoadConfigCollectionV2ForLoadPartitions(ctx, &querypb.LoadPartitionsRequest{
CollectionID: autoscalePrecheckTestCollectionID, ReplicaNumber: 1,
PartitionIDs: []int64{autoscalePrecheckTestPartitionID},
})
} else {
err = server.broadcastAlterLoadConfigCollectionV2ForLoadCollection(ctx, &querypb.LoadCollectionRequest{
CollectionID: autoscalePrecheckTestCollectionID, ReplicaNumber: 1,
})
}
switch scenario {
case "precheck rejection":
var precheckErr *loadResourcePrecheckError
require.ErrorAs(t, err, &precheckErr)
bapi.AssertNotCalled(t, "Broadcast", mock.Anything, mock.Anything)
case "broadcast failure", "metrics and broadcast failure":
require.ErrorIs(t, err, broadcastErr)
case "success", "tiered success", "metrics failure":
require.NoError(t, err)
memoryBefore += 1024
diskBefore += 2048
}
assert.Equal(t, memoryBefore, testutil.ToFloat64(memoryCounter))
assert.Equal(t, diskBefore, testutil.ToFloat64(diskCounter))
if tiered {
cluster.AssertNotCalled(t, "GetMetrics", mock.Anything, mock.Anything, mock.Anything)
}
})
}
}
}
func buildAlterLoadConfigBroadcastResult(collectionID int64) message.BroadcastResultAlterLoadConfigMessageV2 {
controlChannel := "_ctrl_channel"
broadcastMsg := message.NewAlterLoadConfigMessageBuilderV2().
internal/querycoordv2/meta/resource_manager.go
+5
View File
@@ -452,6 +452,8 @@ func (rm *ResourceManager) DropResourceGroup(ctx context.Context, rgName string)
metrics.QueryCoordResourceGroupReplicaTotal.DeletePartialMatch(prometheus.Labels{
metrics.ResourceGroupLabelName: rgName,
})
metrics.QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(rgName)
metrics.QueryCoordLoadDemandDiskBytes.DeleteLabelValues(rgName)
mlog.Info(context.TODO(), "remove resource group",
mlog.String("rgName", rgName),
@@ -1185,6 +1187,9 @@ func (rm *ResourceManager) validateResourceGroupIsDeletable(rgName string) error
// setupInMemResourceGroup setup resource group in memory.
func (rm *ResourceManager) setupInMemResourceGroup(r *ResourceGroup) {
metrics.QueryCoordLoadDemandMemoryBytes.WithLabelValues(r.GetName()).Add(0)
metrics.QueryCoordLoadDemandDiskBytes.WithLabelValues(r.GetName()).Add(0)
// clear old metrics and nodeIDMap entries.
// Use GetAllNodes (bypasses label filter) to ensure all physical nodes are cleaned up,
// even when the RG's label filter has changed.
internal/querycoordv2/meta/resource_manager_test.go
+28
View File
@@ -21,6 +21,7 @@ import (
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
@@ -29,10 +30,12 @@ import (
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/metastore/kv/querycoord"
metastoremocks "github.com/milvus-io/milvus/internal/metastore/mocks"
"github.com/milvus-io/milvus/internal/querycoordv2/params"
"github.com/milvus-io/milvus/internal/querycoordv2/session"
"github.com/milvus-io/milvus/internal/util/sessionutil"
"github.com/milvus-io/milvus/pkg/v3/kv"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/util/etcd"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
@@ -79,6 +82,31 @@ func TestResourceManager(t *testing.T) {
suite.Run(t, new(ResourceManagerSuite))
}
func TestLoadDemandMetricsFollowResourceGroupLifecycle(t *testing.T) {
const rgName = "load-demand-metrics-lifecycle"
metrics.QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(rgName)
metrics.QueryCoordLoadDemandDiskBytes.DeleteLabelValues(rgName)
t.Cleanup(func() {
metrics.QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(rgName)
metrics.QueryCoordLoadDemandDiskBytes.DeleteLabelValues(rgName)
})
catalog := metastoremocks.NewQueryCoordCatalog(t)
catalog.EXPECT().RemoveResourceGroup(mock.Anything, rgName).Return(nil).Once()
nodeMgr := session.NewNodeManager()
manager := NewResourceManager(catalog, nodeMgr)
manager.setupInMemResourceGroup(NewResourceGroup(rgName, newResourceGroupConfig(0, 0), nodeMgr))
require.True(t, metrics.QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(rgName))
require.True(t, metrics.QueryCoordLoadDemandDiskBytes.DeleteLabelValues(rgName))
metrics.QueryCoordLoadDemandMemoryBytes.WithLabelValues(rgName).Add(0)
metrics.QueryCoordLoadDemandDiskBytes.WithLabelValues(rgName).Add(0)
require.NoError(t, manager.DropResourceGroup(context.Background(), rgName))
require.False(t, metrics.QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(rgName))
require.False(t, metrics.QueryCoordLoadDemandDiskBytes.DeleteLabelValues(rgName))
}
func (suite *ResourceManagerSuite) TestValidateConfiguration() {
ctx := suite.ctx
err := suite.manager.validateResourceGroupConfig("rg1", newResourceGroupConfig(0, 0))
internal/querycoordv2/server_test.go
+3
View File
@@ -669,7 +669,10 @@ func (suite *ServerSuite) expectGetRecoverInfo(collection int64) {
})
}
}
suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collection, mock.Anything).Maybe().Return(vChannels, segmentInfos, nil)
suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collection, mock.Anything, mock.Anything).Maybe().Return(vChannels, segmentInfos, nil)
suite.broker.EXPECT().GetRecoveryInfoV2(mock.Anything, collection).Maybe().Return(vChannels, segmentInfos, nil)
suite.broker.EXPECT().GetIndexInfo(mock.Anything, collection, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Maybe().Return(map[int64][]*querypb.FieldIndexInfo{}, nil)
}
func (suite *ServerSuite) expectLoadAndReleasePartitions(querynode *mocks.MockQueryNode) {
internal/querycoordv2/services.go
+2 -2
View File
@@ -273,7 +273,7 @@ func (s *Server) LoadCollection(ctx context.Context, req *querypb.LoadCollection
if err := s.broadcastAlterLoadConfigCollectionV2ForLoadCollection(ctx, req); err != nil {
logger.Warn(ctx, "failed to load collection", mlog.Err(err))
metrics.QueryCoordLoadCount.WithLabelValues(metrics.FailLabel).Inc()
return merr.Status(err), nil
return statusWithLoadResourcePrecheckSuggestion(err), nil
}
logger.Info(ctx, "load collection done")
@@ -346,7 +346,7 @@ func (s *Server) LoadPartitions(ctx context.Context, req *querypb.LoadPartitions
if err := s.broadcastAlterLoadConfigCollectionV2ForLoadPartitions(ctx, req); err != nil {
logger.Warn(ctx, "failed to load partitions", mlog.Err(err))
metrics.QueryCoordLoadCount.WithLabelValues(metrics.FailLabel).Inc()
return merr.Status(err), nil
return statusWithLoadResourcePrecheckSuggestion(err), nil
}
logger.Info(ctx, "load partitions done")
metrics.QueryCoordLoadCount.WithLabelValues(metrics.SuccessLabel).Inc()
internal/querycoordv2/services_test.go
+31 -3
View File
@@ -2120,9 +2120,37 @@ func (suite *ServiceSuite) expectGetRecoverInfo(collection int64) {
})
}
}
suite.broker.EXPECT().
GetRecoveryInfoV2(mock.Anything, collection, mock.Anything, mock.Anything).
Return(vChannels, segmentBinlogs, nil).Maybe()
getRecoveryInfo := func(_ context.Context, _ int64, partitionIDs ...int64) ([]*datapb.VchannelInfo, []*datapb.SegmentInfo, error) {
if len(partitionIDs) == 0 {
return vChannels, segmentBinlogs, nil
}
partitionSet := typeutil.NewSet(partitionIDs...)
filteredSegments := lo.Filter(segmentBinlogs, func(segment *datapb.SegmentInfo, _ int) bool {
return partitionSet.Contain(segment.GetPartitionID())
})
return vChannels, filteredSegments, nil
}
for partitionNum := 0; partitionNum <= len(suite.partitions[collection])+1; partitionNum++ {
partitionArgs := make([]interface{}, 0, partitionNum)
for i := 0; i < partitionNum; i++ {
partitionArgs = append(partitionArgs, mock.Anything)
}
suite.broker.EXPECT().
GetRecoveryInfoV2(mock.Anything, collection, partitionArgs...).
RunAndReturn(getRecoveryInfo).
Maybe()
}
for segmentNum := 0; segmentNum <= len(segmentBinlogs); segmentNum++ {
indexArgs := make([]interface{}, 0, segmentNum)
for i := 0; i < segmentNum; i++ {
indexArgs = append(indexArgs, mock.Anything)
}
suite.broker.EXPECT().
GetIndexInfo(mock.Anything, collection, indexArgs...).
Return(nil, nil).
Maybe()
}
}
func (suite *ServiceSuite) expectLoadMetaRPCs() {
internal/querycoordv2/task/utils.go
+5 -3
View File
@@ -273,7 +273,7 @@ func packUnsubDmChannelRequest(task *ChannelTask, action Action) *querypb.UnsubD
}
// applyCollectionSettings applies collection-level mmap and warmup settings to all fields.
// It combines applyCollectionMmapSetting and applyCollectionWarmupSetting into a single call.
// It combines ApplyCollectionMmapSetting and applyCollectionWarmupSetting into a single call.
func applyCollectionSettings(schema *schemapb.CollectionSchema,
collectionProperties []*commonpb.KeyValuePair,
) *schemapb.CollectionSchema {
@@ -281,7 +281,7 @@ func applyCollectionSettings(schema *schemapb.CollectionSchema,
schemaCloned := typeutil.Clone(schema)
schemaCloned.Properties = mergeCollectionProps(schemaCloned.Properties, collectionProperties)
schemaCloned = applyCollectionMmapSetting(schemaCloned, collectionProperties)
schemaCloned = ApplyCollectionMmapSetting(schemaCloned, collectionProperties)
schemaCloned = applyCollectionWarmupSetting(schemaCloned, collectionProperties)
// Index warmup is normally materialized into each IndexInfo by
@@ -299,7 +299,9 @@ func applyCollectionSettings(schema *schemapb.CollectionSchema,
return schemaCloned
}
func applyCollectionMmapSetting(schema *schemapb.CollectionSchema,
// ApplyCollectionMmapSetting applies collection and struct mmap defaults in place,
// preserving explicit field settings. Callers must pass an owned schema.
func ApplyCollectionMmapSetting(schema *schemapb.CollectionSchema,
collectionProperties []*commonpb.KeyValuePair,
) *schemapb.CollectionSchema {
// field mmap enabled if collection-level mmap enabled or the field mmap enabled
internal/querynodev2/metrics_info.go
+20
View File
@@ -259,6 +259,24 @@ func getSegmentJSON(node *QueryNode, collectionID int64) string {
return string(ret)
}
func getLocalStorageMetrics(loader segments.Loader) *metricsinfo.QueryNodeLocalStorageMetrics {
if loader == nil {
return nil
}
diskUsed, err := loader.GetLocalDiskUsage()
if err != nil || diskUsed < 0 {
return nil
}
capacity := paramtable.Get().QueryNodeCfg.DiskCapacityLimit.GetAsInt64()
if capacity <= 0 {
return nil
}
return &metricsinfo.QueryNodeLocalStorageMetrics{
CapacityBytes: capacity,
UsedBytes: diskUsed,
}
}
// getSystemInfoMetrics returns metrics info of QueryNode
func getSystemInfoMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest, node *QueryNode) (string, error) {
usedMem := hardware.GetUsedMemoryCount()
@@ -297,6 +315,7 @@ func getSystemInfoMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest,
JemallocOverhead: jemallocStats.Overhead,
JemallocSuccess: jemallocStats.Success,
}
localStorage := getLocalStorageMetrics(node.loader)
quotaMetrics, err := getQuotaMetrics(node)
if err != nil {
@@ -325,6 +344,7 @@ func getSystemInfoMetrics(ctx context.Context, req *milvuspb.GetMetricsRequest,
},
QuotaMetrics: quotaMetrics,
CollectionMetrics: collectionMetrics,
LocalStorage: localStorage,
}
metricsinfo.FillDeployMetricsWithEnv(&nodeInfos.SystemInfo)
internal/querynodev2/metrics_info_test.go
+22
View File
@@ -23,6 +23,7 @@ import (
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/distributed/streaming"
@@ -167,3 +168,24 @@ func TestStreamingQuotaMetrics(t *testing.T) {
m = getStreamingQuotaMetrics()
assert.Nil(t, m)
}
func TestGetLocalStorageMetrics(t *testing.T) {
paramtable.Init()
params := paramtable.Get()
require.NoError(t, params.Save(params.QueryNodeCfg.DiskCapacityLimit.Key, "1600"))
t.Cleanup(func() {
params.Reset(params.QueryNodeCfg.DiskCapacityLimit.Key)
})
const gib = int64(1024 * 1024 * 1024)
loader := segments.NewMockLoader(t)
loader.EXPECT().GetLocalDiskUsage().Return(2*gib, nil).Once()
metrics := getLocalStorageMetrics(loader)
require.NotNil(t, metrics)
require.Equal(t, int64(1600)*gib, metrics.CapacityBytes)
require.Equal(t, int64(2)*gib, metrics.UsedBytes)
missingLoader := segments.NewMockLoader(t)
missingLoader.EXPECT().GetLocalDiskUsage().Return(0, errors.New("disk usage unavailable")).Once()
require.Nil(t, getLocalStorageMetrics(missingLoader))
}
internal/querynodev2/segments/load_index_info.go
-66
View File
@@ -1,66 +0,0 @@
// 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 segments
/*
#cgo pkg-config: milvus_core
#include <stdlib.h>
#include "segcore/load_index_c.h"
#include "common/binary_set_c.h"
*/
import "C"
import (
"context"
"unsafe"
)
// LoadIndexInfo is a wrapper of the underlying C-structure C.CLoadIndexInfo
type LoadIndexInfo struct {
cLoadIndexInfo C.CLoadIndexInfo
}
// newLoadIndexInfo creates a CLoadIndexInfo on the current DynamicPool worker.
func newLoadIndexInfo(ctx context.Context) (*LoadIndexInfo, error) {
var cLoadIndexInfo C.CLoadIndexInfo
status := C.NewLoadIndexInfo(&cLoadIndexInfo)
if err := HandleCStatus(ctx, &status, "NewLoadIndexInfo failed"); err != nil {
return nil, err
}
return &LoadIndexInfo{cLoadIndexInfo: cLoadIndexInfo}, nil
}
// deleteLoadIndexInfo deletes a CLoadIndexInfo on the current DynamicPool worker.
func deleteLoadIndexInfo(info *LoadIndexInfo) {
C.DeleteLoadIndexInfo(info.cLoadIndexInfo)
}
func (li *LoadIndexInfo) appendLoadIndexInfo(ctx context.Context, marshaled []byte) error {
status := C.FinishLoadIndexInfo(li.cLoadIndexInfo, (*C.uint8_t)(unsafe.Pointer(&marshaled[0])), (C.uint64_t)(len(marshaled)))
return HandleCStatus(ctx, &status, "FinishLoadIndexInfo failed")
}
func (li *LoadIndexInfo) setShard(shard string) {
if shard == "" {
return
}
cShard := C.CString(shard)
defer C.free(unsafe.Pointer(cShard))
C.SetLoadIndexInfoShard(li.cLoadIndexInfo, cShard)
}
internal/querynodev2/segments/mock_loader.go
+55
View File
@@ -72,6 +72,61 @@ func (_c *MockLoader_GetChunkManager_Call) RunAndReturn(run func() storage.Chunk
return _c
}
// GetLocalDiskUsage provides a mock function with no fields
func (_m *MockLoader) GetLocalDiskUsage() (int64, error) {
ret := _m.Called()
if len(ret) == 0 {
panic("no return value specified for GetLocalDiskUsage")
}
var r0 int64
var r1 error
if rf, ok := ret.Get(0).(func() (int64, error)); ok {
return rf()
}
if rf, ok := ret.Get(0).(func() int64); ok {
r0 = rf()
} else {
r0 = ret.Get(0).(int64)
}
if rf, ok := ret.Get(1).(func() error); ok {
r1 = rf()
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockLoader_GetLocalDiskUsage_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetLocalDiskUsage'
type MockLoader_GetLocalDiskUsage_Call struct {
*mock.Call
}
// GetLocalDiskUsage is a helper method to define mock.On call
func (_e *MockLoader_Expecter) GetLocalDiskUsage() *MockLoader_GetLocalDiskUsage_Call {
return &MockLoader_GetLocalDiskUsage_Call{Call: _e.mock.On("GetLocalDiskUsage")}
}
func (_c *MockLoader_GetLocalDiskUsage_Call) Run(run func()) *MockLoader_GetLocalDiskUsage_Call {
_c.Call.Run(func(args mock.Arguments) {
run()
})
return _c
}
func (_c *MockLoader_GetLocalDiskUsage_Call) Return(_a0 int64, _a1 error) *MockLoader_GetLocalDiskUsage_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockLoader_GetLocalDiskUsage_Call) RunAndReturn(run func() (int64, error)) *MockLoader_GetLocalDiskUsage_Call {
_c.Call.Return(run)
return _c
}
// Load provides a mock function with given fields: ctx, collectionID, segmentType, version, _a4
func (_m *MockLoader) Load(ctx context.Context, collectionID int64, segmentType commonpb.SegmentState, version int64, _a4 ...*querypb.SegmentLoadInfo) ([]Segment, error) {
_va := make([]interface{}, len(_a4))
internal/querynodev2/segments/segment.go
+4 -94
View File
@@ -47,19 +47,14 @@ import (
"github.com/milvus-io/milvus/internal/querynodev2/segments/state"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
"github.com/milvus-io/milvus/internal/util/segcore"
"github.com/milvus-io/milvus/internal/util/vecindexmgr"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/metrics"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/proto/cgopb"
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/proto/segcorepb"
"github.com/milvus-io/milvus/pkg/v3/util/contextutil"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/indexparams"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/metautil"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
@@ -222,6 +217,8 @@ func (s *baseSegment) compactLoadInfoForRuntime() {
loadInfo := s.LoadInfo()
usage, err := estimateLogicalResourceUsageOfSegment(s.collection.Schema(), loadInfo, resourceEstimateFactor{
deltaDataExpansionFactor: paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
jsonKeyStatsExpansionFactor: paramtable.Get().QueryNodeCfg.JSONKeyStatsExpansionFactor.GetAsFloat(),
textIndexExpansionFactor: paramtable.Get().QueryNodeCfg.TextIndexExpansionFactor.GetAsFloat(),
TieredEvictionEnabled: paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool(),
TieredEvictableMemoryCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableMemoryCacheRatio.GetAsFloat(),
TieredEvictableDiskCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableDiskCacheRatio.GetAsFloat(),
@@ -397,6 +394,8 @@ func (s *baseSegment) ResourceUsageEstimate() ResourceUsage {
usage, err := estimateLogicalResourceUsageOfSegment(s.collection.Schema(), s.LoadInfo(), resourceEstimateFactor{
deltaDataExpansionFactor: paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
jsonKeyStatsExpansionFactor: paramtable.Get().QueryNodeCfg.JSONKeyStatsExpansionFactor.GetAsFloat(),
textIndexExpansionFactor: paramtable.Get().QueryNodeCfg.TextIndexExpansionFactor.GetAsFloat(),
TieredEvictionEnabled: paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool(),
TieredEvictableMemoryCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableMemoryCacheRatio.GetAsFloat(),
TieredEvictableDiskCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableDiskCacheRatio.GetAsFloat(),
@@ -1147,95 +1146,6 @@ func (s *LocalSegment) LoadDeltaData(ctx context.Context, deltaData *storage.Del
return nil
}
func GetCLoadInfoWithFunc(ctx context.Context,
fieldSchema *schemapb.FieldSchema,
loadInfo *querypb.SegmentLoadInfo,
indexInfo *querypb.FieldIndexInfo,
f func(c *LoadIndexInfo) error,
) error {
indexParams := funcutil.KeyValuePair2Map(indexInfo.IndexParams)
// as Knowhere reports error if encounter an unknown param, we need to delete it
delete(indexParams, common.MmapEnabledKey)
// some build params also exist in indexParams, which are useless during loading process
if vecindexmgr.GetVecIndexMgrInstance().IsDiskANN(indexParams["index_type"]) {
if err := indexparams.SetDiskIndexLoadParams(paramtable.Get(), indexParams, indexInfo.GetNumRows()); err != nil {
return err
}
}
// set whether enable offset cache for bitmap index
if indexParams["index_type"] == indexparamcheck.IndexBitmap {
indexparams.SetBitmapIndexLoadParams(paramtable.Get(), indexParams)
}
if err := indexparams.AppendPrepareLoadParams(paramtable.Get(), indexParams); err != nil {
return err
}
enableMmap := isIndexMmapEnable(fieldSchema, indexInfo)
// Add warmup policy to index_params if not already present
// C++ will pass it to Knowhere for index loading
if existingWarmup, exists := indexParams[common.WarmupKey]; exists {
mlog.Info(ctx, "warmup policy already in index params (from QueryCoord)",
mlog.FieldSegmentID(loadInfo.GetSegmentID()),
mlog.FieldFieldID(indexInfo.GetFieldID()),
mlog.String("warmup", existingWarmup))
} else {
warmupPolicy := getIndexWarmupPolicy(fieldSchema, indexInfo)
mlog.Info(ctx, "warmup policy from getIndexWarmupPolicy",
mlog.FieldSegmentID(loadInfo.GetSegmentID()),
mlog.FieldFieldID(indexInfo.GetFieldID()),
mlog.String("warmup", warmupPolicy))
if warmupPolicy != "" {
indexParams[common.WarmupKey] = warmupPolicy
}
}
// Pass DataCoord-built index file paths through; QueryNode should not
// attach v0/v1 path layout semantics to the read path.
indexInfoProto := &cgopb.LoadIndexInfo{
CollectionID: loadInfo.GetCollectionID(),
PartitionID: loadInfo.GetPartitionID(),
SegmentID: loadInfo.GetSegmentID(),
Field: fieldSchema,
EnableMmap: enableMmap,
IndexID: indexInfo.GetIndexID(),
IndexBuildID: indexInfo.GetBuildID(),
IndexVersion: indexInfo.GetIndexVersion(),
IndexParams: indexParams,
IndexFiles: indexInfo.GetIndexFilePaths(),
IndexEngineVersion: indexInfo.GetCurrentIndexVersion(),
IndexFileSize: indexInfo.GetIndexSize(),
NumRows: indexInfo.GetNumRows(),
CurrentScalarIndexVersion: indexInfo.GetCurrentScalarIndexVersion(),
IndexStorePathVersion: indexInfo.GetIndexStorePathVersion(),
}
marshaled, err := proto.Marshal(indexInfoProto)
if err != nil {
return err
}
// Keep the complete CLoadIndexInfo lifecycle in one DynamicPool task. The
// callback runs on that worker and must call C APIs directly instead of
// submitting another task to DynamicPool.
_, err = GetDynamicPool().Submit(func() (any, error) {
loadIndexInfo, err := newLoadIndexInfo(ctx)
if err != nil {
return nil, err
}
defer deleteLoadIndexInfo(loadIndexInfo)
if err := loadIndexInfo.appendLoadIndexInfo(ctx, marshaled); err != nil {
mlog.Warn(ctx, "fail to append load index info", mlog.Err(err))
return nil, err
}
loadIndexInfo.setShard(loadInfo.GetInsertChannel())
return nil, f(loadIndexInfo)
}).Await()
return err
}
func (s *LocalSegment) syncFieldJSONStatsFromLoadInfo(ctx context.Context, loadInfo *querypb.SegmentLoadInfo) {
jsonStatsInfo := make(map[int64]*querypb.JsonStatsInfo)
if !paramtable.Get().CommonCfg.EnabledJSONKeyStats.GetAsBool() {
internal/querynodev2/segments/segment_loader.go
+43 -564
View File
@@ -49,6 +49,7 @@ import (
"github.com/milvus-io/milvus/internal/storagecommon"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
"github.com/milvus-io/milvus/internal/util/segcore/loadresource"
"github.com/milvus-io/milvus/internal/util/vecindexmgr"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/metrics"
@@ -62,7 +63,6 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/logutil"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/metautil"
"github.com/milvus-io/milvus/pkg/v3/util/metric"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
"github.com/milvus-io/milvus/pkg/v3/util/timerecord"
@@ -91,30 +91,15 @@ type Loader interface {
// GetChunkManager returns the chunk manager for remote storage access.
GetChunkManager() storage.ChunkManager
// GetLocalDiskUsage returns the cached size of the local storage directory.
GetLocalDiskUsage() (int64, error)
// ReopenSegments update segment data according to new load info.
ReopenSegments(ctx context.Context,
loadInfos []*querypb.SegmentLoadInfo,
) error
}
type ResourceEstimate struct {
MaxMemoryCost uint64
MaxDiskCost uint64
FinalMemoryCost uint64
FinalDiskCost uint64
HasRawData bool
}
func GetResourceEstimate(estimate *C.LoadResourceRequest) ResourceEstimate {
return ResourceEstimate{
MaxMemoryCost: uint64(estimate.max_memory_cost),
MaxDiskCost: uint64(estimate.max_disk_cost),
FinalMemoryCost: uint64(estimate.final_memory_cost),
FinalDiskCost: uint64(estimate.final_disk_cost),
HasRawData: bool(estimate.has_raw_data),
}
}
type requestResourceResult struct {
Resource LoadResource
LogicalResource LoadResource
@@ -161,22 +146,6 @@ type resourceEstimateFactor struct {
externalRawDataFactor float64
}
func estimateTantivyValidityBitmapBytes(numRows int64) uint64 {
if numRows <= 0 {
return 0
}
// TextIndexStats does not carry the null count, so admission reserves the
// full word-aligned bitmap conservatively. All-valid indexes charge zero
// actual bitmap bytes after loading.
const (
bitsPerWord = uint64(64)
bytesPerWord = uint64(8)
)
words := (uint64(numRows)-1)/bitsPerWord + 1
return words * bytesPerWord
}
func NewLoader(
ctx context.Context,
manager *Manager,
@@ -239,6 +208,10 @@ type segmentLoader struct {
var _ Loader = (*segmentLoader)(nil)
func (loader *segmentLoader) GetLocalDiskUsage() (int64, error) {
return loader.duf.GetDiskUsage()
}
func (loader *segmentLoader) Load(ctx context.Context,
collectionID int64,
segmentType SegmentType,
@@ -1720,6 +1693,7 @@ func (loader *segmentLoader) checkLogicalSegmentSize(ctx context.Context, segmen
// so we need to estimate the final resource usage of the segments
finalFactor := resourceEstimateFactor{
deltaDataExpansionFactor: paramtable.Get().QueryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
jsonKeyStatsExpansionFactor: paramtable.Get().QueryNodeCfg.JSONKeyStatsExpansionFactor.GetAsFloat(),
textIndexExpansionFactor: paramtable.Get().QueryNodeCfg.TextIndexExpansionFactor.GetAsFloat(),
TieredEvictionEnabled: paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.GetAsBool(),
TieredEvictableMemoryCacheRatio: paramtable.Get().QueryNodeCfg.TieredEvictableMemoryCacheRatio.GetAsFloat(),
@@ -1926,243 +1900,37 @@ func (loader *segmentLoader) checkLoadingResource(
return nil
}
// resolveSegmentEstimateLogs returns the raw and delta metadata consumed by the
// existing estimator. Internal Storage V3 descriptors are adapted to pathless
// FieldBinlogs; all other SegmentLoadInfo metadata remains on the original proto.
func resolveSegmentEstimateLogs(schema *schemapb.CollectionSchema, loadInfo *querypb.SegmentLoadInfo) ([]*datapb.FieldBinlog, []*datapb.FieldBinlog) {
binlogs := loadInfo.GetBinlogPaths()
deltalogs := loadInfo.GetDeltalogs()
if len(binlogs) > 0 {
return binlogs, deltalogs
}
if loadInfo.GetStorageVersion() != storage.StorageV3 || loadInfo.GetManifestPath() == "" || typeutil.IsExternalCollection(schema) {
return binlogs, deltalogs
}
stats := loadInfo.GetStats()
resource := stats.GetLoadResource()
if resource == nil {
return binlogs, deltalogs
}
groups := resource.GetColumnGroups()
binlogs = make([]*datapb.FieldBinlog, 0, len(groups))
for _, group := range groups {
binlogs = append(binlogs, &datapb.FieldBinlog{
FieldID: group.GetGroupId(),
ChildFields: append([]int64(nil), group.GetFieldIds()...),
Binlogs: []*datapb.Binlog{{
EntriesNum: loadInfo.GetNumOfRows(),
MemorySize: group.GetMemorySize(),
}},
})
}
deltalogs = nil
if deltaSize := stats.GetDeltaBinlogSize(); deltaSize > 0 {
deltalogs = []*datapb.FieldBinlog{{
Binlogs: []*datapb.Binlog{{MemorySize: deltaSize}},
}}
}
return binlogs, deltalogs
}
// this function is used to estimate the logical resource usage of a segment, which should only be used when tiered eviction is enabled
// the result is the final resource usage of the segment inevictable part plus the final usage of evictable part with cache ratio applied
// TODO: the inevictable part is not correct, since we cannot know the final resource usage of interim index and default-value column before loading,
// current they are ignored, but we should consider them in the future
func estimateLogicalResourceUsageOfSegment(schema *schemapb.CollectionSchema, loadInfo *querypb.SegmentLoadInfo, multiplyFactor resourceEstimateFactor) (usage *ResourceUsage, err error) {
var segmentInevictableMemorySize, segmentInevictableDiskSize uint64
var segmentEvictableMemorySize, segmentEvictableDiskSize uint64
binlogs, deltalogs := resolveSegmentEstimateLogs(schema, loadInfo)
id2Binlogs := lo.SliceToMap(binlogs, func(fieldBinlog *datapb.FieldBinlog) (int64, *datapb.FieldBinlog) {
return fieldBinlog.GetFieldID(), fieldBinlog
options := loadresource.DefaultSegmentFinalEstimateOptions()
options.DeltaDataExpansionFactor = multiplyFactor.deltaDataExpansionFactor
options.JSONKeyStatsExpansionFactor = multiplyFactor.jsonKeyStatsExpansionFactor
options.TextIndexExpansionFactor = multiplyFactor.textIndexExpansionFactor
options.TieredEvictionEnabled = multiplyFactor.TieredEvictionEnabled
options.TieredEvictableMemoryCacheRatio = multiplyFactor.TieredEvictableMemoryCacheRatio
options.TieredEvictableDiskCacheRatio = multiplyFactor.TieredEvictableDiskCacheRatio
estimate, err := loadresource.EstimateSegmentFinalResource(context.Background(), schema, loadInfo, options, func(fn func() error) error {
_, err := GetDynamicPool().Submit(func() (any, error) {
return nil, fn()
}).Await()
return err
})
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
if err != nil {
mlog.Warn(context.TODO(), "failed to create schema helper", mlog.String("name", schema.GetName()), mlog.Err(err))
return nil, err
}
ctx := context.TODO()
// PART 1: calculate logical resource usage of indexes
for _, fieldIndexInfo := range loadInfo.IndexInfos {
fieldID := fieldIndexInfo.GetFieldID()
if len(fieldIndexInfo.GetIndexFilePaths()) > 0 {
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
return nil, err
}
isVectorType := typeutil.IsVectorType(fieldSchema.GetDataType())
var estimateResult ResourceEstimate
err = GetCLoadInfoWithFunc(ctx, fieldSchema, loadInfo, fieldIndexInfo, func(c *LoadIndexInfo) error {
var loadResourceRequest C.LoadResourceRequest
status := C.EstimateLoadIndexResource(c.cLoadIndexInfo, &loadResourceRequest)
if err := HandleCStatus(ctx, &status, "failed to estimate load index resource"); err != nil {
return err
}
estimateResult = GetResourceEstimate(&loadResourceRequest)
return nil
})
if err != nil {
return nil, merr.Wrapf(err, "failed to estimate logical resource usage of index, collection %d, segment %d, indexBuildID %d",
loadInfo.GetCollectionID(),
loadInfo.GetSegmentID(),
fieldIndexInfo.GetBuildID())
}
segmentEvictableMemorySize += estimateResult.FinalMemoryCost
segmentEvictableDiskSize += estimateResult.FinalDiskCost
// could skip binlog or
// could be missing for new field or storage v2 group 0
if estimateResult.HasRawData &&
!paramtable.Get().QueryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool() {
delete(id2Binlogs, fieldID)
continue
}
// BM25 only checks vector datatype
// scalar index does not have metrics type key
if !isVectorType {
continue
}
metricType, err := funcutil.GetAttrByKeyFromRepeatedKV(common.MetricTypeKey, fieldIndexInfo.IndexParams)
if err != nil {
return nil, merr.Wrapf(err, "failed to estimate logical resource usage of index, metric type not found, collection %d, segment %d, indexBuildID %d",
loadInfo.GetCollectionID(),
loadInfo.GetSegmentID(),
fieldIndexInfo.GetBuildID())
}
// skip raw data for BM25 index
if metricType == metric.BM25 {
delete(id2Binlogs, fieldID)
}
}
}
// PART 2: calculate logical resource usage of binlogs
for fieldID, fieldBinlog := range id2Binlogs {
fieldIDs := fieldBinlog.GetChildFields()
// legacy default split
if len(fieldIDs) == 0 {
fieldIDs = []int64{fieldID}
}
binlogSize := uint64(getBinlogDataMemorySize(fieldBinlog))
var supportInterimIndexDataType bool
var containsTimestampField bool
var doubleMemoryDataField bool
var legacyNilSchema bool
mmapEnabled := true
isVectorType := true
for _, fieldID := range fieldIDs {
// get field schema from fieldID
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
mlog.Info(ctx, "skip binlog for dropped field", mlog.FieldFieldID(fieldID), mlog.String("name", schema.GetName()))
continue
}
// missing mapping, shall be "0" group for storage v2
if fieldSchema == nil {
legacyNilSchema = true
break
}
supportInterimIndexDataType = supportInterimIndexDataType || SupportInterimIndexDataType(fieldSchema.GetDataType())
isVectorType = isVectorType && typeutil.IsVectorType(fieldSchema.GetDataType())
// constainSystemField = constainSystemField || common.IsSystemField(fieldSchema.GetFieldID())
mmapEnabled = mmapEnabled && isDataMmapEnable(fieldSchema)
containsTimestampField = containsTimestampField || DoubleMemorySystemField(fieldSchema.GetFieldID())
doubleMemoryDataField = doubleMemoryDataField || DoubleMemoryDataType(fieldSchema.GetDataType())
}
// TODO: add default-value column's resource usage to inevictable part
// TODO: add interim index's resource usage to inevictable part
if legacyNilSchema {
segmentEvictableMemorySize += binlogSize
continue
}
// timestamp field double in InsertRecord & TimestampIndex
if containsTimestampField {
timestampSize := lo.SumBy(fieldBinlog.GetBinlogs(), func(binlog *datapb.Binlog) int64 {
return binlog.GetEntriesNum() * 4
})
segmentInevictableMemorySize += 2 * uint64(timestampSize)
}
if isVectorType {
mmapVectorField := paramtable.Get().QueryNodeCfg.MmapVectorField.GetAsBool()
if mmapVectorField {
segmentEvictableDiskSize += binlogSize
} else {
segmentEvictableMemorySize += binlogSize
}
} else if !mmapEnabled {
segmentEvictableMemorySize += binlogSize
if doubleMemoryDataField {
segmentEvictableMemorySize += binlogSize
}
} else {
segmentEvictableDiskSize += binlogSize
}
}
// PART 3: calculate logical resource usage of stats data
for _, fieldBinlog := range loadInfo.Statslogs {
segmentInevictableMemorySize += uint64(getBinlogDataMemorySize(fieldBinlog))
}
// PART 4: calculate logical resource usage of delete data
for _, fieldBinlog := range deltalogs {
// MemorySize of filedBinlog is the actual size in memory, so the expansionFactor
// should be 1, in most cases.
expansionFactor := float64(1)
memSize := getBinlogDataMemorySize(fieldBinlog)
// Note: If MemorySize == DiskSize, it means the segment comes from Milvus 2.3,
// MemorySize is actually compressed DiskSize of deltalog, so we'll fallback to use
// deltaExpansionFactor to compromise the compression ratio.
if memSize == getBinlogDataDiskSize(fieldBinlog) {
expansionFactor = multiplyFactor.deltaDataExpansionFactor
}
segmentInevictableMemorySize += uint64(float64(memSize) * expansionFactor)
}
// PART 5: calculate logical resource usage of text index stats data
// Text match indexes are evictable (support_eviction=true in caching layer).
// Text match index mmap is driven by scalar_field_enable_mmap (same as raw scalar data).
textIndexMmapEnable := paramtable.Get().QueryNodeCfg.MmapScalarField.GetAsBool()
validityBitmapBytes := estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
for _, textStats := range loadInfo.GetTextStatsLogs() {
indexFileBytes := uint64(float64(textStats.GetMemorySize()) * multiplyFactor.textIndexExpansionFactor)
segmentEvictableMemorySize += validityBitmapBytes
if textIndexMmapEnable {
segmentEvictableDiskSize += indexFileBytes
} else {
segmentEvictableMemorySize += indexFileBytes
}
}
mlog.Debug(context.TODO(), "estimate logical resoure usage result",
mlog.Debug(context.TODO(), "estimate logical resource usage result",
mlog.Int64("segmentID", loadInfo.GetSegmentID()),
mlog.Uint64("segmentInevictableMemorySize", segmentInevictableMemorySize),
mlog.Uint64("segmentEvictableMemorySize", segmentEvictableMemorySize),
mlog.Uint64("segmentInevictableDiskSize", segmentInevictableDiskSize),
mlog.Uint64("segmentEvictableDiskSize", segmentEvictableDiskSize),
mlog.Uint64("memorySize", estimate.MemoryBytes),
mlog.Uint64("diskSize", estimate.DiskBytes),
)
return &ResourceUsage{
MemorySize: segmentInevictableMemorySize + uint64(float64(segmentEvictableMemorySize)*multiplyFactor.TieredEvictableMemoryCacheRatio),
DiskSize: segmentInevictableDiskSize + uint64(float64(segmentEvictableDiskSize)*multiplyFactor.TieredEvictableDiskCacheRatio),
MemorySize: estimate.MemoryBytes,
DiskSize: estimate.DiskBytes,
}, nil
}
@@ -2172,322 +1940,33 @@ func estimateLogicalResourceUsageOfSegment(schema *schemapb.CollectionSchema, lo
// which should be a subset of the segment inevictable part
// - when tiered eviction is disabled, the result is the max resource usage of both the segment evictable and inevictable part
func estimateLoadingResourceUsageOfSegment(schema *schemapb.CollectionSchema, loadInfo *querypb.SegmentLoadInfo, multiplyFactor resourceEstimateFactor) (usage *ResourceUsage, err error) {
var segMemoryLoadingSize, segDiskLoadingSize uint64
var indexMemorySize uint64
var mmapFieldCount int
var fieldGpuMemorySize []uint64
options := loadresource.DefaultSegmentLoadingEstimateOptions()
options.DeltaDataExpansionFactor = multiplyFactor.deltaDataExpansionFactor
options.JSONKeyStatsExpansionFactor = multiplyFactor.jsonKeyStatsExpansionFactor
options.TextIndexExpansionFactor = multiplyFactor.textIndexExpansionFactor
options.TieredEvictionEnabled = multiplyFactor.TieredEvictionEnabled
options.EnableInterimSegmentIndex = multiplyFactor.EnableInterminSegmentIndex
options.TempSegmentIndexFactor = multiplyFactor.tempSegmentIndexFactor
options.ExternalRawDataFactor = multiplyFactor.externalRawDataFactor
binlogs, deltalogs := resolveSegmentEstimateLogs(schema, loadInfo)
id2Binlogs := lo.SliceToMap(binlogs, func(fieldBinlog *datapb.FieldBinlog) (int64, *datapb.FieldBinlog) {
return fieldBinlog.GetFieldID(), fieldBinlog
estimate, err := loadresource.EstimateSegmentLoadingResource(context.Background(), schema, loadInfo, options, func(fn func() error) error {
_, err := GetDynamicPool().Submit(func() (any, error) {
return nil, fn()
}).Await()
return err
})
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
if err != nil {
mlog.Warn(context.TODO(), "failed to create schema helper", mlog.String("name", schema.GetName()), mlog.Err(err))
return nil, err
}
indexedFields := make(map[int64]struct{})
ctx := context.Background()
// PART 1: calculate size of indexes
for _, fieldIndexInfo := range loadInfo.IndexInfos {
fieldID := fieldIndexInfo.GetFieldID()
if len(fieldIndexInfo.GetIndexFilePaths()) > 0 {
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
// field might have been dropped, skip its index
mlog.Info(ctx, "skip index for dropped field", mlog.FieldFieldID(fieldID), mlog.String("name", schema.GetName()))
continue
}
indexedFields[fieldID] = struct{}{}
isVectorType := typeutil.IsVectorType(fieldSchema.GetDataType())
var estimateResult ResourceEstimate
err = GetCLoadInfoWithFunc(ctx, fieldSchema, loadInfo, fieldIndexInfo, func(c *LoadIndexInfo) error {
var loadResourceRequest C.LoadResourceRequest
status := C.EstimateLoadIndexResource(c.cLoadIndexInfo, &loadResourceRequest)
if err := HandleCStatus(ctx, &status, "failed to estimate load index resource"); err != nil {
return err
}
estimateResult = GetResourceEstimate(&loadResourceRequest)
return nil
})
if err != nil {
return nil, merr.Wrapf(err, "failed to estimate loading resource usage of index, collection %d, segment %d, indexBuildID %d",
loadInfo.GetCollectionID(),
loadInfo.GetSegmentID(),
fieldIndexInfo.GetBuildID())
}
if !multiplyFactor.TieredEvictionEnabled {
indexMemorySize += estimateResult.MaxMemoryCost
segDiskLoadingSize += estimateResult.MaxDiskCost
}
if gpuIndexRequiresGpu(fieldIndexInfo.IndexParams) {
fieldGpuMemorySize = append(fieldGpuMemorySize, estimateResult.MaxMemoryCost)
}
// could skip binlog or
// could be missing for new field or storage v2 group 0
if estimateResult.HasRawData &&
!paramtable.Get().QueryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool() {
delete(id2Binlogs, fieldID)
continue
}
// BM25 only checks vector datatype
// scalar index does not have metrics type key
if !isVectorType {
continue
}
metricType, err := funcutil.GetAttrByKeyFromRepeatedKV(common.MetricTypeKey, fieldIndexInfo.IndexParams)
if err != nil {
return nil, merr.Wrapf(err, "failed to estimate loading resource usage of index, metric type not found, collection %d, segment %d, indexBuildID %d",
loadInfo.GetCollectionID(),
loadInfo.GetSegmentID(),
fieldIndexInfo.GetBuildID())
}
// skip raw data for BM25 index
if metricType == metric.BM25 {
delete(id2Binlogs, fieldID)
}
}
}
// PART 2: calculate size of binlogs
for fieldID, fieldBinlog := range id2Binlogs {
fieldIDs := fieldBinlog.GetChildFields()
// legacy default split
if len(fieldIDs) == 0 {
fieldIDs = []int64{fieldID}
}
binlogSize := uint64(getBinlogDataMemorySize(fieldBinlog))
var supportInterimIndexDataType bool
var containsTimestampField bool
var doubleMomoryDataField bool
var legacyNilSchema bool
mmapEnabled := true
isVectorType := true
hasIndex := true
for _, fieldID := range fieldIDs {
// get field schema from fieldID
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
// field might have been dropped, skip it and continue processing
// other fields in the same column group
mlog.Info(ctx, "skip binlog for dropped field", mlog.FieldFieldID(fieldID), mlog.String("name", schema.GetName()))
continue
}
if _, ok := indexedFields[fieldID]; !ok {
hasIndex = false
}
// missing mapping, shall be "0" group for storage v2
if fieldSchema == nil {
if !multiplyFactor.TieredEvictionEnabled {
segMemoryLoadingSize += binlogSize
}
legacyNilSchema = true
break
}
supportInterimIndexDataType = supportInterimIndexDataType || SupportInterimIndexDataType(fieldSchema.GetDataType())
isVectorType = isVectorType && typeutil.IsVectorType(fieldSchema.GetDataType())
mmapEnabled = mmapEnabled && isDataMmapEnable(fieldSchema)
containsTimestampField = containsTimestampField || DoubleMemorySystemField(fieldSchema.GetFieldID())
doubleMomoryDataField = doubleMomoryDataField || DoubleMemoryDataType(fieldSchema.GetDataType())
}
// legacy v2 segment without children
if legacyNilSchema {
continue
}
if !hasIndex {
if !multiplyFactor.TieredEvictionEnabled {
interimIndexEnable := multiplyFactor.EnableInterminSegmentIndex && !isGrowingMmapEnable() && supportInterimIndexDataType
if interimIndexEnable {
segMemoryLoadingSize += uint64(float64(binlogSize) * multiplyFactor.tempSegmentIndexFactor)
}
}
}
if isVectorType {
mmapVectorField := paramtable.Get().QueryNodeCfg.MmapVectorField.GetAsBool()
if mmapVectorField {
if !multiplyFactor.TieredEvictionEnabled {
segDiskLoadingSize += binlogSize
}
} else {
if !multiplyFactor.TieredEvictionEnabled {
segMemoryLoadingSize += binlogSize
}
}
continue
}
// timestamp field double in InsertRecord & TimestampIndex
if containsTimestampField {
timestampSize := lo.SumBy(fieldBinlog.GetBinlogs(), func(binlog *datapb.Binlog) int64 {
return binlog.GetEntriesNum() * 4
})
segMemoryLoadingSize += 2 * uint64(timestampSize)
}
if !mmapEnabled {
if !multiplyFactor.TieredEvictionEnabled {
segMemoryLoadingSize += binlogSize
if doubleMomoryDataField {
segMemoryLoadingSize += binlogSize
}
}
} else {
if !multiplyFactor.TieredEvictionEnabled {
segDiskLoadingSize += uint64(getBinlogDataMemorySize(fieldBinlog))
}
}
}
// PART 2.5: external segment adjustments
//
// External segments carry pre-computed MemorySize in fake binlogs (from
// DataNode Take sampling). Adjust the memory estimate for two external-
// specific behaviors:
// 1. Non-lazy path without tiered eviction: apply externalRawDataFactor
// to cover the peak transient memory during download + decompress +
// Arrow deserialize (normal packed segments do not have this peak
// because their binlogs are already in Arrow IPC format). When tiered
// eviction is enabled, the caching layer reserves this loading
// overhead, so Go must not reserve the same raw-data margin again.
// 2. Full-lazy path (all external fields warmup=disable): no eager
// load, so subtract the raw data size that PART 2 added.
// Also propagate EstimatedBytesPerRow to the C++ ManifestGroupTranslator
// so the tiered-cache layer sizes chunks correctly.
if typeutil.IsExternalCollection(schema) && loadInfo.GetNumOfRows() > 0 {
var fakeBinlogMemSize int64
for _, fb := range binlogs {
fakeBinlogMemSize += getBinlogDataMemorySize(fb)
}
loadInfo.EstimatedBytesPerRow = fakeBinlogMemSize / loadInfo.GetNumOfRows()
if isExternalCollectionLazyLoad(schema) {
// Full-lazy → zero eager load. Undo PART 2's rawSize addition.
// Safety factor does not apply: no peak to cover.
if segMemoryLoadingSize >= uint64(fakeBinlogMemSize) {
segMemoryLoadingSize -= uint64(fakeBinlogMemSize)
} else {
segMemoryLoadingSize = 0
}
} else if !multiplyFactor.TieredEvictionEnabled {
if factor := multiplyFactor.externalRawDataFactor; factor > 1.0 {
// Non-lazy → add peak margin on top of rawSize that PART 2 added.
segMemoryLoadingSize += uint64(float64(fakeBinlogMemSize) * (factor - 1.0))
}
}
}
// PART 3: calculate size of stats data
// stats data isn't managed by the caching layer, so its size should always be included,
// regardless of the tiered eviction value
for _, fieldBinlog := range loadInfo.Statslogs {
segMemoryLoadingSize += uint64(getBinlogDataMemorySize(fieldBinlog))
}
// PART 4: calculate size of delete data
// delete data isn't managed by the caching layer, so its size should always be included,
// regardless of the tiered eviction value
for _, fieldBinlog := range deltalogs {
// MemorySize of filedBinlog is the actual size in memory, but we should also consider
// the memcpy from golang to cpp side, so the expansionFactor is set to 2.
expansionFactor := float64(2)
memSize := getBinlogDataMemorySize(fieldBinlog)
// Note: If MemorySize == DiskSize, it means the segment comes from Milvus 2.3,
// MemorySize is actually compressed DiskSize of deltalog, so we'll fallback to use
// deltaExpansionFactor to compromise the compression ratio.
if memSize == getBinlogDataDiskSize(fieldBinlog) {
expansionFactor = multiplyFactor.deltaDataExpansionFactor
}
segMemoryLoadingSize += uint64(float64(memSize) * expansionFactor)
}
// PART 5: calculate size of json key stats data
jsonStatsMmapEnable := paramtable.Get().QueryNodeCfg.MmapJSONStats.GetAsBool()
for _, jsonKeyStats := range loadInfo.GetJsonKeyStatsLogs() {
if jsonStatsMmapEnable {
if !multiplyFactor.TieredEvictionEnabled {
segDiskLoadingSize += uint64(float64(jsonKeyStats.GetMemorySize()) * multiplyFactor.jsonKeyStatsExpansionFactor)
}
} else {
if !multiplyFactor.TieredEvictionEnabled {
segMemoryLoadingSize += uint64(float64(jsonKeyStats.GetMemorySize()) * multiplyFactor.jsonKeyStatsExpansionFactor)
}
}
}
// per struct memory size, used to keep mapping between row id and element id
var structArrayOffsetsSize uint64
// PART 6: calculate size of struct array offsets
// The memory size is 4 * row_count + 4 * total_element_count
// We cannot easily get the element count, so we estimate it by the row count * 10
rowCount := uint64(loadInfo.GetNumOfRows())
for range len(schema.GetStructArrayFields()) {
structArrayOffsetsSize += 4*rowCount + 4*rowCount*10
}
// PART 7: calculate size of text index stats data
// text index data is managed by the caching layer when tiered eviction is enabled,
// so it only needs to be included when tiered eviction is disabled.
// Text match index mmap is driven by scalar_field_enable_mmap (same as raw scalar data).
// memory_size is the sum of uploaded Tantivy index files, including sparse null sidecars.
// The materialized word-aligned validity bitmap is separate heap memory, and
// textIndexExpansionFactor applies only to the index file bytes.
textIndexMmapEnable := paramtable.Get().QueryNodeCfg.MmapScalarField.GetAsBool()
validityBitmapBytes := estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
for _, textStats := range loadInfo.GetTextStatsLogs() {
if multiplyFactor.TieredEvictionEnabled {
continue
}
indexFileBytes := uint64(float64(textStats.GetMemorySize()) * multiplyFactor.textIndexExpansionFactor)
segMemoryLoadingSize += validityBitmapBytes
if textIndexMmapEnable {
segDiskLoadingSize += indexFileBytes
} else {
segMemoryLoadingSize += indexFileBytes
}
}
return &ResourceUsage{
MemorySize: segMemoryLoadingSize + indexMemorySize + structArrayOffsetsSize,
DiskSize: segDiskLoadingSize,
MmapFieldCount: mmapFieldCount,
FieldGpuMemorySize: fieldGpuMemorySize,
MemorySize: estimate.MemoryBytes,
DiskSize: estimate.DiskBytes,
MmapFieldCount: estimate.MmapFieldCount,
FieldGpuMemorySize: estimate.FieldGPUMemoryBytes,
}, nil
}
func DoubleMemoryDataType(dataType schemapb.DataType) bool {
return dataType == schemapb.DataType_String ||
dataType == schemapb.DataType_VarChar ||
dataType == schemapb.DataType_JSON
}
func DoubleMemorySystemField(fieldID int64) bool {
return fieldID == common.TimeStampField
}
func SupportInterimIndexDataType(dataType schemapb.DataType) bool {
return dataType == schemapb.DataType_FloatVector ||
dataType == schemapb.DataType_SparseFloatVector ||
dataType == schemapb.DataType_Float16Vector ||
dataType == schemapb.DataType_BFloat16Vector
}
// prepareIndexLoadParams injects QueryNode-local index load parameters into each
// index's IndexParams in place. These params (e.g. DISKANN num_load_thread) are
// derived from local QueryNode resources/config and are never persisted in the
internal/querynodev2/segments/segment_loader_test.go
+88 -21
View File
@@ -37,6 +37,7 @@ import (
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/storagev2/packed"
"github.com/milvus-io/milvus/internal/util/initcore"
"github.com/milvus-io/milvus/internal/util/segcore/loadresource"
"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"
@@ -50,6 +51,27 @@ import (
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
func TestGetLocalDiskUsage(t *testing.T) {
duf := &diskUsageFetcher{
usage: atomic.NewInt64(0),
err: atomic.NewError(nil),
}
loader := &segmentLoader{duf: duf}
usage, err := loader.GetLocalDiskUsage()
assert.NoError(t, err)
assert.Zero(t, usage)
duf.usage.Store(100)
loader.committedResource.DiskSize = 20
usage, err = loader.GetLocalDiskUsage()
assert.NoError(t, err)
assert.Equal(t, int64(100), usage)
duf.err.Store(merr.WrapErrServiceInternalMsg("disk usage unavailable"))
_, err = loader.GetLocalDiskUsage()
assert.Error(t, err)
}
func TestSeparateLoadInfoV2LegacyStatsRoot(t *testing.T) {
paramtable.Init()
params := paramtable.Get()
@@ -2066,13 +2088,11 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) baseLoadInfo(textStats map[int
}
}
func (suite *SegmentLoaderTextIndexEstimateSuite) TestTantivyValidityBitmapBytesWordAligned() {
suite.EqualValues(0, estimateTantivyValidityBitmapBytes(-1))
suite.EqualValues(0, estimateTantivyValidityBitmapBytes(0))
suite.EqualValues(8, estimateTantivyValidityBitmapBytes(1))
suite.EqualValues(8, estimateTantivyValidityBitmapBytes(64))
suite.EqualValues(16, estimateTantivyValidityBitmapBytes(65))
suite.EqualValues(24, estimateTantivyValidityBitmapBytes(129))
func expectedTantivyValidityBitmapBytes(numRows int64) uint64 {
if numRows <= 0 {
return 0
}
return ((uint64(numRows)-1)/64 + 1) * 8
}
func (suite *SegmentLoaderTextIndexEstimateSuite) TestLoadingEstimate_NonMmap_NoTieredEviction() {
@@ -2092,7 +2112,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLoadingEstimate_NonMmap_No
}
usage, err := estimateLoadingResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
expectedMemorySize := uint64(textIndexSize) + estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
expectedMemorySize := uint64(textIndexSize) + expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
suite.EqualValues(expectedMemorySize, usage.MemorySize, "non-mmap text index files and validity bitmap must be counted in memory")
suite.EqualValues(0, usage.DiskSize, "non-mmap text index must not be counted in disk")
}
@@ -2114,7 +2134,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLoadingEstimate_Mmap_NoTie
}
usage, err := estimateLoadingResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
suite.EqualValues(estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows()), usage.MemorySize, "mmap text index validity bitmap must be counted in memory")
suite.EqualValues(expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows()), usage.MemorySize, "mmap text index validity bitmap must be counted in memory")
suite.EqualValues(textIndexSize, usage.DiskSize, "mmap text index must be counted in disk")
}
@@ -2159,7 +2179,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLoadingEstimate_MultipleTe
}
usage, err := estimateLoadingResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
validityBitmapBytes := estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
validityBitmapBytes := expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
suite.EqualValues(uint64(size1+size2)+2*validityBitmapBytes, usage.MemorySize, "each text field must include its word-aligned validity bitmap")
}
@@ -2180,7 +2200,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLoadingEstimate_ExpansionF
}
usage, err := estimateLoadingResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
expected := uint64(float64(textIndexSize)*expansionFactor) + estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
expected := uint64(float64(textIndexSize)*expansionFactor) + expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
suite.EqualValues(expected, usage.MemorySize, "expansion factor must apply only to text index file bytes")
}
@@ -2242,7 +2262,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLogicalEstimate_NonMmap_Ev
}
usage, err := estimateLogicalResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
expectedMemorySize := uint64(textIndexSize) + estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
expectedMemorySize := uint64(textIndexSize) + expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
suite.EqualValues(expectedMemorySize, usage.MemorySize, "non-mmap text index files and validity bitmap must be in evictable memory")
suite.EqualValues(0, usage.DiskSize)
}
@@ -2265,7 +2285,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLogicalEstimate_Mmap_Evict
}
usage, err := estimateLogicalResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
suite.EqualValues(estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows()), usage.MemorySize, "mmap text index validity bitmap must be in evictable memory")
suite.EqualValues(expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows()), usage.MemorySize, "mmap text index validity bitmap must be in evictable memory")
suite.EqualValues(textIndexSize, usage.DiskSize, "mmap text index must be in evictable disk")
}
@@ -2289,7 +2309,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLogicalEstimate_CacheRatio
}
usage, err := estimateLogicalResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
evictableMemorySize := uint64(textIndexSize) + estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
evictableMemorySize := uint64(textIndexSize) + expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
expected := uint64(float64(evictableMemorySize) * cacheRatio)
suite.EqualValues(expected, usage.MemorySize, "cache ratio must be applied to evictable text index files and validity bitmap")
}
@@ -2333,7 +2353,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLogicalEstimate_MultipleTe
}
usage, err := estimateLogicalResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
validityBitmapBytes := estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
validityBitmapBytes := expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
suite.EqualValues(uint64(size1+size2)+2*validityBitmapBytes, usage.MemorySize, "each text field must include its word-aligned validity bitmap in logical estimate")
}
@@ -2357,7 +2377,7 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLogicalEstimate_DiskCacheR
}
usage, err := estimateLogicalResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
suite.EqualValues(estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows()), usage.MemorySize, "mmap text index validity bitmap must remain in memory")
suite.EqualValues(expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows()), usage.MemorySize, "mmap text index validity bitmap must remain in memory")
expected := uint64(float64(textIndexSize) * diskCacheRatio)
suite.EqualValues(expected, usage.DiskSize, "disk cache ratio must be applied to mmap text index")
}
@@ -2381,10 +2401,57 @@ func (suite *SegmentLoaderTextIndexEstimateSuite) TestLogicalEstimate_ExpansionF
}
usage, err := estimateLogicalResourceUsageOfSegment(suite.schema, loadInfo, factor)
suite.NoError(err)
expected := uint64(float64(textIndexSize)*expansionFactor) + estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
expected := uint64(float64(textIndexSize)*expansionFactor) + expectedTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
suite.EqualValues(expected, usage.MemorySize)
}
func TestCheckLogicalSegmentSizeUsesJSONKeyStatsExpansionFactor(t *testing.T) {
paramtable.Init()
params := paramtable.Get()
params.Save(params.QueryNodeCfg.TieredEvictionEnabled.Key, "true")
defer params.Reset(params.QueryNodeCfg.TieredEvictionEnabled.Key)
params.Save(params.QueryNodeCfg.TieredEvictableMemoryCacheRatio.Key, "1.0")
defer params.Reset(params.QueryNodeCfg.TieredEvictableMemoryCacheRatio.Key)
params.Save(params.QueryNodeCfg.MmapJSONStats.Key, "false")
defer params.Reset(params.QueryNodeCfg.MmapJSONStats.Key)
params.Save(params.QueryNodeCfg.JSONKeyStatsExpansionFactor.Key, "3.0")
defer params.Reset(params.QueryNodeCfg.JSONKeyStatsExpansionFactor.Key)
const collectionID = int64(10)
collectionManager := NewMockCollectionManager(t)
segmentManager := NewMockSegmentManager(t)
loader := &segmentLoader{
manager: &Manager{
Collection: collectionManager,
Segment: segmentManager,
},
}
collectionManager.EXPECT().
Get(collectionID).
Return(NewCollectionWithoutSegcoreForTest(collectionID, &schemapb.CollectionSchema{
Name: "test_json_stats_estimate",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "id", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
},
}))
segmentManager.EXPECT().GetLogicalResource().Return(ResourceUsage{}).Twice()
memoryBytes, diskBytes, err := loader.checkLogicalSegmentSize(context.Background(), []*querypb.SegmentLoadInfo{
{
CollectionID: collectionID,
SegmentID: 20,
NumOfRows: 1,
JsonKeyStatsLogs: map[int64]*datapb.JsonKeyStats{
101: {FieldID: 101, MemorySize: 100},
},
},
}, 1024)
assert.NoError(t, err)
assert.EqualValues(t, 300, memoryBytes)
assert.Zero(t, diskBytes)
}
func TestSeparateLoadInfoV2_ExternalFieldIndexNotSkipped(t *testing.T) {
// Verifies that indexes on external fields are NOT skipped in separateLoadInfoV2.
// Previously, external field indexes were filtered out, preventing index loading for external tables.
@@ -2845,7 +2912,7 @@ func TestResolveSegmentEstimateLogs(t *testing.T) {
}
loadInfo := newLoadInfo()
binlogs, deltalogs := resolveSegmentEstimateLogs(schema, loadInfo)
binlogs, deltalogs := loadresource.ResolveSegmentEstimateLogs(schema, loadInfo)
assert.Empty(t, loadInfo.GetBinlogPaths())
assert.Len(t, binlogs, 1)
assert.Equal(t, int64(0), binlogs[0].GetFieldID())
@@ -2885,7 +2952,7 @@ func TestResolveSegmentEstimateLogs(t *testing.T) {
assert.Equal(t, manualUsage, adaptedUsage)
loadInfo.GetStats().InsertBinlogSize = 99
binlogs, _ = resolveSegmentEstimateLogs(schema, loadInfo)
binlogs, _ = loadresource.ResolveSegmentEstimateLogs(schema, loadInfo)
assert.Len(t, binlogs, 1)
for name, mutate := range map[string]func(*querypb.SegmentLoadInfo, *schemapb.CollectionSchema){
@@ -2908,7 +2975,7 @@ func TestResolveSegmentEstimateLogs(t *testing.T) {
info.BinlogPaths = []*datapb.FieldBinlog{{FieldID: 99}}
testSchema := newSchema()
mutate(info, testSchema)
gotBinlogs, gotDeltalogs := resolveSegmentEstimateLogs(testSchema, info)
gotBinlogs, gotDeltalogs := loadresource.ResolveSegmentEstimateLogs(testSchema, info)
assert.Equal(t, info.GetBinlogPaths(), gotBinlogs)
assert.Equal(t, info.GetDeltalogs(), gotDeltalogs)
})
@@ -2919,7 +2986,7 @@ func TestResolveSegmentEstimateLogs(t *testing.T) {
info.Stats.InsertBinlogSize = 0
info.Stats.LoadResource.ColumnGroups = nil
gotBinlogs, gotDeltalogs := resolveSegmentEstimateLogs(schema, info)
gotBinlogs, gotDeltalogs := loadresource.ResolveSegmentEstimateLogs(schema, info)
assert.Empty(t, gotBinlogs)
assert.Len(t, gotDeltalogs, 1)
assert.EqualValues(t, 50, gotDeltalogs[0].GetBinlogs()[0].GetMemorySize())
internal/querynodev2/segments/segment_test.go
+16 -2
View File
@@ -242,8 +242,16 @@ func TestCompactLoadInfoForRuntimeCachesResourceUsage(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.MmapScalarField.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.MmapScalarField.Key)
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.MmapJSONStats.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.MmapJSONStats.Key)
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.Key, "false")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredEvictionEnabled.Key)
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredEvictableMemoryCacheRatio.Key, "1.0")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredEvictableMemoryCacheRatio.Key)
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.JSONKeyStatsExpansionFactor.Key, "3.0")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.JSONKeyStatsExpansionFactor.Key)
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TextIndexExpansionFactor.Key, "2.0")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TextIndexExpansionFactor.Key)
schema := &schemapb.CollectionSchema{
Name: "test",
@@ -269,6 +277,12 @@ func TestCompactLoadInfoForRuntimeCachesResourceUsage(t *testing.T) {
}},
},
},
JsonKeyStatsLogs: map[int64]*datapb.JsonKeyStats{
101: {FieldID: 101, MemorySize: 100},
},
TextStatsLogs: map[int64]*datapb.TextIndexStats{
101: {FieldID: 101, MemorySize: 200},
},
}
segment := &baseSegment{
collection: NewTestCollection(loadInfo.GetCollectionID(), querypb.LoadType_LoadCollection, schema),
@@ -283,9 +297,9 @@ func TestCompactLoadInfoForRuntimeCachesResourceUsage(t *testing.T) {
assert.Nil(t, segment.LoadInfo().GetStats())
cached := segment.resourceUsageCache.Load()
if assert.NotNil(t, cached) {
assert.EqualValues(t, 4096, cached.MemorySize)
assert.EqualValues(t, 4796, cached.MemorySize)
}
assert.EqualValues(t, 4096, segment.ResourceUsageEstimate().MemorySize)
assert.EqualValues(t, 4796, segment.ResourceUsageEstimate().MemorySize)
}
func (suite *SegmentSuite) TestSyncFieldJSONStatsFromLoadInfo() {
internal/querynodev2/segments/utils.go
-37
View File
@@ -11,8 +11,6 @@ import (
"github.com/milvus-io/milvus/internal/querycoordv2/params"
"github.com/milvus-io/milvus/internal/querynodev2/segments/metricsutil"
"github.com/milvus-io/milvus/internal/storage"
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
"github.com/milvus-io/milvus/internal/util/vecindexmgr"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/mlog"
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
@@ -258,25 +256,6 @@ func getFieldSchema(schema *schemapb.CollectionSchema, fieldID int64) (*schemapb
return nil, merr.WrapErrFieldNotFound(fieldID, "not in schema")
}
func isIndexMmapEnable(fieldSchema *schemapb.FieldSchema, indexInfo *querypb.FieldIndexInfo) bool {
enableMmap, exist := common.IsMmapIndexEnabled(indexInfo.IndexParams...)
// fast path for returning disabled, need to perform index type check for enabled case
if exist && !enableMmap {
return enableMmap
}
indexType := common.GetIndexType(indexInfo.IndexParams)
var indexSupportMmap bool
// var defaultEnableMmap bool
if typeutil.IsVectorType(fieldSchema.GetDataType()) {
indexSupportMmap = vecindexmgr.GetVecIndexMgrInstance().IsMMapSupported(indexType)
enableMmap = params.Params.QueryNodeCfg.MmapVectorIndex.GetAsBool() || enableMmap
} else {
indexSupportMmap = indexparamcheck.IsScalarMmapIndex(indexType)
enableMmap = params.Params.QueryNodeCfg.MmapScalarIndex.GetAsBool() || enableMmap
}
return indexSupportMmap && enableMmap
}
// Except accepting whether the raw data is loaded in mmap or not, it also affects the stats index such as
// text match index and json key stats index.
func isDataMmapEnable(fieldSchema *schemapb.FieldSchema) bool {
@@ -310,22 +289,6 @@ func getFieldWarmupPolicy(fieldSchema *schemapb.FieldSchema) string {
return params.Params.QueryNodeCfg.TieredWarmupScalarField.GetValue()
}
// getIndexWarmupPolicy returns the warmup policy for index loading.
// Priority: index params (propagated from collection-level by QueryCoord, including autoWarmupForNonPKIsolationCollection) > global config
func getIndexWarmupPolicy(fieldSchema *schemapb.FieldSchema, indexInfo *querypb.FieldIndexInfo) string {
// Check index params (collection-level warmup.scalarIndex/warmup.vectorIndex
// and autoWarmupForNonPKIsolationCollection are propagated to index params by QueryCoord)
policy, exist := common.GetWarmupPolicy(indexInfo.IndexParams...)
if exist {
return policy
}
// Fall back to global config
if typeutil.IsVectorType(fieldSchema.GetDataType()) {
return params.Params.QueryNodeCfg.TieredWarmupVectorIndex.GetValue()
}
return params.Params.QueryNodeCfg.TieredWarmupScalarIndex.GetValue()
}
// getScalarDataWarmupPolicy returns the warmup policy for scalar data, but also include json key stats and text match.
// Priority: field TypeParams (propagated from collection-level by QueryCoord, including autoWarmupForNonPKIsolationCollection) > global config
func getScalarDataWarmupPolicy(fieldSchema *schemapb.FieldSchema) string {
internal/querynodev2/segments/utils_test.go
-145
View File
@@ -155,71 +155,6 @@ func TestGetFieldSchema(t *testing.T) {
})
}
func TestIsIndexMmapEnable(t *testing.T) {
paramtable.Init()
t.Run("mmap index param exist", func(t *testing.T) {
enable := isIndexMmapEnable(&schemapb.FieldSchema{}, &querypb.FieldIndexInfo{
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.MmapEnabledKey,
Value: "false",
},
},
})
assert.False(t, enable)
})
t.Run("mmap vector index param not exist", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.MmapVectorIndex.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.MmapVectorIndex.Key)
enable := isIndexMmapEnable(&schemapb.FieldSchema{
DataType: schemapb.DataType_FloatVector,
}, &querypb.FieldIndexInfo{
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "IVF_FLAT",
},
},
})
assert.True(t, enable)
})
t.Run("mmap scalar index param not exist", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.MmapScalarIndex.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.MmapScalarIndex.Key)
enable := isIndexMmapEnable(&schemapb.FieldSchema{
DataType: schemapb.DataType_String,
}, &querypb.FieldIndexInfo{
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "INVERTED",
},
},
})
assert.True(t, enable)
})
t.Run("mmap scalar index param not supported", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.MmapScalarIndex.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.MmapScalarIndex.Key)
enable := isIndexMmapEnable(&schemapb.FieldSchema{
DataType: schemapb.DataType_String,
}, &querypb.FieldIndexInfo{
IndexParams: []*commonpb.KeyValuePair{
{
Key: common.IndexTypeKey,
Value: "STL_SORT",
},
{Key: common.MmapEnabledKey, Value: "true"},
},
})
assert.False(t, enable)
})
}
func TestIsDataMmmapEnable(t *testing.T) {
paramtable.Init()
@@ -334,86 +269,6 @@ func TestGetFieldWarmupPolicy(t *testing.T) {
})
}
func TestGetIndexWarmupPolicy(t *testing.T) {
paramtable.Init()
t.Run("index params has warmup", func(t *testing.T) {
policy := getIndexWarmupPolicy(
&schemapb.FieldSchema{DataType: schemapb.DataType_String},
&querypb.FieldIndexInfo{
IndexParams: []*commonpb.KeyValuePair{
{Key: common.WarmupKey, Value: common.WarmupSync},
},
},
)
assert.Equal(t, common.WarmupSync, policy)
})
t.Run("index params warmup propagated from collection level", func(t *testing.T) {
policy := getIndexWarmupPolicy(
&schemapb.FieldSchema{DataType: schemapb.DataType_String},
&querypb.FieldIndexInfo{
IndexParams: []*commonpb.KeyValuePair{
{Key: common.WarmupKey, Value: common.WarmupDisable},
},
},
)
assert.Equal(t, common.WarmupDisable, policy)
})
t.Run("async policy in index params is returned", func(t *testing.T) {
policy := getIndexWarmupPolicy(
&schemapb.FieldSchema{DataType: schemapb.DataType_String},
&querypb.FieldIndexInfo{
IndexParams: []*commonpb.KeyValuePair{
{Key: common.WarmupKey, Value: common.WarmupAsync},
},
},
)
assert.Equal(t, common.WarmupAsync, policy)
})
t.Run("fallback to global config for scalar index", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredWarmupScalarIndex.Key, common.WarmupSync)
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredWarmupScalarIndex.Key)
policy := getIndexWarmupPolicy(
&schemapb.FieldSchema{DataType: schemapb.DataType_String},
&querypb.FieldIndexInfo{},
)
assert.Equal(t, common.WarmupSync, policy)
})
t.Run("fallback to global config for scalar index async", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredWarmupScalarIndex.Key, common.WarmupAsync)
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredWarmupScalarIndex.Key)
policy := getIndexWarmupPolicy(
&schemapb.FieldSchema{DataType: schemapb.DataType_String},
&querypb.FieldIndexInfo{},
)
assert.Equal(t, common.WarmupAsync, policy)
})
t.Run("fallback to global config for vector index", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredWarmupVectorIndex.Key, common.WarmupDisable)
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredWarmupVectorIndex.Key)
policy := getIndexWarmupPolicy(
&schemapb.FieldSchema{DataType: schemapb.DataType_FloatVector},
&querypb.FieldIndexInfo{},
)
assert.Equal(t, common.WarmupDisable, policy)
})
t.Run("fallback to global config for vector index async", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredWarmupVectorIndex.Key, common.WarmupAsync)
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredWarmupVectorIndex.Key)
policy := getIndexWarmupPolicy(
&schemapb.FieldSchema{DataType: schemapb.DataType_FloatVector},
&querypb.FieldIndexInfo{},
)
assert.Equal(t, common.WarmupAsync, policy)
})
}
func TestGetScalarDataWarmupPolicy(t *testing.T) {
paramtable.Init()
internal/util/segcore/loadresource/index.go
+196
View File
@@ -0,0 +1,196 @@
// 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 loadresource
/*
#cgo pkg-config: milvus_core
#include <stdlib.h>
#include "segcore/load_index_c.h"
*/
import "C"
import (
"context"
"unsafe"
"github.com/cockroachdb/errors"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
"github.com/milvus-io/milvus/internal/util/indexparamcheck"
"github.com/milvus-io/milvus/internal/util/vecindexmgr"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/proto/cgopb"
"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/indexparams"
"github.com/milvus-io/milvus/pkg/v3/util/merr"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type IndexResourceUsage struct {
MaxMemoryBytes uint64
MaxDiskBytes uint64
FinalMemoryBytes uint64
FinalDiskBytes uint64
HasRawData bool
}
type Runner func(func() error) error
func EstimateIndexLoadResource(ctx context.Context, field *schemapb.FieldSchema, segment *querypb.SegmentLoadInfo, index *querypb.FieldIndexInfo) (IndexResourceUsage, error) {
return EstimateIndexLoadResourceWithRunner(ctx, field, segment, index, nil)
}
func EstimateIndexLoadResourceWithRunner(_ context.Context, field *schemapb.FieldSchema, segment *querypb.SegmentLoadInfo, index *querypb.FieldIndexInfo, runner Runner) (IndexResourceUsage, error) {
if field == nil {
return IndexResourceUsage{}, merr.WrapErrServiceInternalMsg("field schema is nil")
}
if segment == nil {
return IndexResourceUsage{}, merr.WrapErrServiceInternalMsg("segment load info is nil")
}
if index == nil {
return IndexResourceUsage{}, merr.WrapErrServiceInternalMsg("field index info is nil")
}
info, err := buildCLoadIndexInfo(field, segment, index)
if err != nil {
return IndexResourceUsage{}, err
}
marshaled, err := proto.Marshal(info)
if err != nil {
return IndexResourceUsage{}, err
}
if len(marshaled) == 0 {
return IndexResourceUsage{}, merr.WrapErrServiceInternalMsg("serialized load index info is empty")
}
var estimate C.LoadResourceRequest
var status C.CStatus
if err := run(runner, func() error {
status = C.EstimateLoadIndexResourceFromSerializedInfo(
(*C.uint8_t)(unsafe.Pointer(&marshaled[0])),
(C.uint64_t)(len(marshaled)),
&estimate,
)
return nil
}); err != nil {
return IndexResourceUsage{}, err
}
if err := consumeCStatus(&status); err != nil {
return IndexResourceUsage{}, errors.Wrap(err, "EstimateLoadIndexResourceFromSerializedInfo failed")
}
return indexResourceUsageFromC(&estimate), nil
}
func run(runner Runner, fn func() error) error {
if runner == nil {
return fn()
}
return runner(fn)
}
func consumeCStatus(status *C.CStatus) error {
if status == nil || status.error_code == 0 {
return nil
}
errorCode := status.error_code
errorMsg := C.GoString(status.error_msg)
C.free(unsafe.Pointer(status.error_msg))
return merr.SegcoreError(int32(errorCode), errorMsg)
}
func buildCLoadIndexInfo(field *schemapb.FieldSchema, segment *querypb.SegmentLoadInfo, index *querypb.FieldIndexInfo) (*cgopb.LoadIndexInfo, error) {
indexParams := funcutil.KeyValuePair2Map(index.GetIndexParams())
delete(indexParams, common.MmapEnabledKey)
indexType := indexParams[common.IndexTypeKey]
if indexType == "" {
return nil, merr.WrapErrServiceInternalMsg("index type is empty")
}
if vecindexmgr.GetVecIndexMgrInstance().IsDiskANN(indexType) {
if err := indexparams.SetDiskIndexLoadParams(paramtable.Get(), indexParams, index.GetNumRows()); err != nil {
return nil, err
}
}
if indexType == indexparamcheck.IndexBitmap {
indexparams.SetBitmapIndexLoadParams(paramtable.Get(), indexParams)
}
if err := indexparams.AppendPrepareLoadParams(paramtable.Get(), indexParams); err != nil {
return nil, err
}
if _, ok := indexParams[common.WarmupKey]; !ok {
if warmupPolicy := getIndexWarmupPolicy(field, index); warmupPolicy != "" {
indexParams[common.WarmupKey] = warmupPolicy
}
}
return &cgopb.LoadIndexInfo{
CollectionID: segment.GetCollectionID(),
PartitionID: segment.GetPartitionID(),
SegmentID: segment.GetSegmentID(),
Field: field,
EnableMmap: isIndexMmapEnable(field, index),
IndexID: index.GetIndexID(),
IndexBuildID: index.GetBuildID(),
IndexVersion: index.GetIndexVersion(),
IndexParams: indexParams,
IndexFiles: index.GetIndexFilePaths(),
IndexEngineVersion: index.GetCurrentIndexVersion(),
IndexFileSize: index.GetIndexSize(),
NumRows: index.GetNumRows(),
CurrentScalarIndexVersion: index.GetCurrentScalarIndexVersion(),
IndexStorePathVersion: index.GetIndexStorePathVersion(),
}, nil
}
func isIndexMmapEnable(field *schemapb.FieldSchema, index *querypb.FieldIndexInfo) bool {
enableMmap, exist := common.IsMmapIndexEnabled(index.GetIndexParams()...)
if exist && !enableMmap {
return false
}
indexType := common.GetIndexType(index.GetIndexParams())
if typeutil.IsVectorType(field.GetDataType()) {
return vecindexmgr.GetVecIndexMgrInstance().IsMMapSupported(indexType) &&
(paramtable.Get().QueryNodeCfg.MmapVectorIndex.GetAsBool() || enableMmap)
}
return indexparamcheck.IsScalarMmapIndex(indexType) &&
(paramtable.Get().QueryNodeCfg.MmapScalarIndex.GetAsBool() || enableMmap)
}
func getIndexWarmupPolicy(field *schemapb.FieldSchema, index *querypb.FieldIndexInfo) string {
if policy, exist := common.GetWarmupPolicy(index.GetIndexParams()...); exist {
return policy
}
if typeutil.IsVectorType(field.GetDataType()) {
return paramtable.Get().QueryNodeCfg.TieredWarmupVectorIndex.GetValue()
}
return paramtable.Get().QueryNodeCfg.TieredWarmupScalarIndex.GetValue()
}
func indexResourceUsageFromC(estimate *C.LoadResourceRequest) IndexResourceUsage {
return IndexResourceUsage{
MaxMemoryBytes: uint64(estimate.max_memory_cost),
MaxDiskBytes: uint64(estimate.max_disk_cost),
FinalMemoryBytes: uint64(estimate.final_memory_cost),
FinalDiskBytes: uint64(estimate.final_disk_cost),
HasRawData: bool(estimate.has_raw_data),
}
}
internal/util/segcore/loadresource/index_test.go
+269
View File
@@ -0,0 +1,269 @@
// 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 loadresource
import (
"context"
"testing"
"github.com/stretchr/testify/require"
"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/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
func TestEstimateIndexLoadResourceDoesNotReadIndexFiles(t *testing.T) {
field := &schemapb.FieldSchema{
FieldID: 100,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "128"},
},
}
segment := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 1000,
}
index := &querypb.FieldIndexInfo{
FieldID: 100,
IndexID: 10,
BuildID: 11,
IndexSize: 1024 * 1024,
NumRows: 1000,
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "HNSW"},
{Key: common.MetricTypeKey, Value: "L2"},
},
IndexFilePaths: []string{"/path/that/must/not/exist/index_file"},
}
usage, err := EstimateIndexLoadResource(context.Background(), field, segment, index)
require.NoError(t, err)
require.NotZero(t, usage.FinalMemoryBytes)
require.True(t, usage.HasRawData)
}
func TestEstimateIndexLoadResourceWithRunner(t *testing.T) {
field := &schemapb.FieldSchema{
FieldID: 100,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "128"},
},
}
segment := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 1000,
}
index := &querypb.FieldIndexInfo{
FieldID: 100,
IndexID: 10,
BuildID: 11,
IndexSize: 1024 * 1024,
NumRows: 1000,
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "HNSW"},
{Key: common.MetricTypeKey, Value: "L2"},
},
}
calls := 0
runner := func(fn func() error) error {
calls++
return fn()
}
usage, err := EstimateIndexLoadResourceWithRunner(context.Background(), field, segment, index, runner)
require.NoError(t, err)
require.NotZero(t, usage.FinalMemoryBytes)
require.Equal(t, 1, calls)
}
func TestBuildCLoadIndexInfoPreservesIndexNumRows(t *testing.T) {
field := &schemapb.FieldSchema{
FieldID: 100,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "128"},
},
}
segment := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 1000,
}
index := &querypb.FieldIndexInfo{
FieldID: 100,
IndexID: 10,
BuildID: 11,
IndexSize: 1024 * 1024,
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "HNSW"},
{Key: common.MetricTypeKey, Value: "L2"},
},
}
info, err := buildCLoadIndexInfo(field, segment, index)
require.NoError(t, err)
require.Zero(t, info.GetNumRows())
}
func TestBuildCLoadIndexInfoPreservesIndexStorePathVersion(t *testing.T) {
field := &schemapb.FieldSchema{
FieldID: 100,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "128"},
},
}
segment := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 1000,
}
index := &querypb.FieldIndexInfo{
FieldID: 100,
IndexID: 10,
BuildID: 11,
IndexSize: 1024 * 1024,
IndexStorePathVersion: indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_COLLECTION_ROOTED,
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "HNSW"},
{Key: common.MetricTypeKey, Value: "L2"},
},
}
info, err := buildCLoadIndexInfo(field, segment, index)
require.NoError(t, err)
require.Equal(t, indexpb.IndexStorePathVersion_INDEX_STORE_PATH_VERSION_COLLECTION_ROOTED, info.GetIndexStorePathVersion())
}
func TestBuildCLoadIndexInfoResolvesMmap(t *testing.T) {
paramtable.Init()
t.Run("explicit disable takes precedence", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.MmapVectorIndex.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.MmapVectorIndex.Key)
info, err := buildCLoadIndexInfo(
&schemapb.FieldSchema{DataType: schemapb.DataType_FloatVector},
&querypb.SegmentLoadInfo{},
&querypb.FieldIndexInfo{IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "IVF_FLAT"},
{Key: common.MmapEnabledKey, Value: "false"},
}},
)
require.NoError(t, err)
require.False(t, info.GetEnableMmap())
})
t.Run("global defaults apply to supported indexes", func(t *testing.T) {
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.MmapScalarIndex.Key, "true")
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.MmapScalarIndex.Key)
info, err := buildCLoadIndexInfo(
&schemapb.FieldSchema{DataType: schemapb.DataType_String},
&querypb.SegmentLoadInfo{},
&querypb.FieldIndexInfo{IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "INVERTED"},
}},
)
require.NoError(t, err)
require.True(t, info.GetEnableMmap())
})
}
func TestBuildCLoadIndexInfoResolvesWarmup(t *testing.T) {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.TieredWarmupScalarIndex.Key, common.WarmupAsync)
defer paramtable.Get().Reset(paramtable.Get().QueryNodeCfg.TieredWarmupScalarIndex.Key)
t.Run("index setting takes precedence", func(t *testing.T) {
info, err := buildCLoadIndexInfo(
&schemapb.FieldSchema{DataType: schemapb.DataType_String},
&querypb.SegmentLoadInfo{},
&querypb.FieldIndexInfo{IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "INVERTED"},
{Key: common.WarmupKey, Value: common.WarmupDisable},
}},
)
require.NoError(t, err)
require.Equal(t, common.WarmupDisable, info.GetIndexParams()[common.WarmupKey])
})
t.Run("global setting is the fallback", func(t *testing.T) {
info, err := buildCLoadIndexInfo(
&schemapb.FieldSchema{DataType: schemapb.DataType_String},
&querypb.SegmentLoadInfo{},
&querypb.FieldIndexInfo{IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "INVERTED"},
}},
)
require.NoError(t, err)
require.Equal(t, common.WarmupAsync, info.GetIndexParams()[common.WarmupKey])
})
}
func TestEstimateIndexLoadResourceReturnsErrorWithoutIndexType(t *testing.T) {
field := &schemapb.FieldSchema{
FieldID: 100,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "128"},
},
}
segment := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 1000,
}
index := &querypb.FieldIndexInfo{
FieldID: 100,
IndexID: 10,
BuildID: 11,
IndexSize: 1024 * 1024,
NumRows: 1000,
}
usage, err := EstimateIndexLoadResource(context.Background(), field, segment, index)
require.Error(t, err)
require.Zero(t, usage)
}
internal/util/segcore/loadresource/segment.go
+778
View File
@@ -0,0 +1,778 @@
// 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 loadresource
import (
"context"
"math"
"strconv"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"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/internal/storage"
"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/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/funcutil"
"github.com/milvus-io/milvus/pkg/v3/util/indexparams"
"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"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
type SegmentResourceUsage struct {
MemoryBytes uint64
DiskBytes uint64
MmapFieldCount int
FieldGPUMemoryBytes []uint64
}
type SegmentFinalEstimateOptions struct {
MmapVectorField bool
MmapScalarField bool
MmapJSONStats bool
DeltaDataExpansionFactor float64
JSONKeyStatsExpansionFactor float64
TextIndexExpansionFactor float64
TieredEvictionEnabled bool
TieredEvictableMemoryCacheRatio float64
TieredEvictableDiskCacheRatio float64
PreferFieldDataWhenIndexHasRawData bool
}
func DefaultSegmentFinalEstimateOptions() SegmentFinalEstimateOptions {
queryNodeCfg := &paramtable.Get().QueryNodeCfg
return SegmentFinalEstimateOptions{
MmapVectorField: queryNodeCfg.MmapVectorField.GetAsBool(),
MmapScalarField: queryNodeCfg.MmapScalarField.GetAsBool(),
MmapJSONStats: queryNodeCfg.MmapJSONStats.GetAsBool(),
DeltaDataExpansionFactor: queryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
JSONKeyStatsExpansionFactor: queryNodeCfg.JSONKeyStatsExpansionFactor.GetAsFloat(),
TextIndexExpansionFactor: queryNodeCfg.TextIndexExpansionFactor.GetAsFloat(),
TieredEvictionEnabled: queryNodeCfg.TieredEvictionEnabled.GetAsBool(),
TieredEvictableMemoryCacheRatio: queryNodeCfg.TieredEvictableMemoryCacheRatio.GetAsFloat(),
TieredEvictableDiskCacheRatio: queryNodeCfg.TieredEvictableDiskCacheRatio.GetAsFloat(),
PreferFieldDataWhenIndexHasRawData: queryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool(),
}
}
type SegmentLoadingEstimateOptions struct {
MmapVectorField bool
MmapScalarField bool
MmapJSONStats bool
DeltaDataExpansionFactor float64
JSONKeyStatsExpansionFactor float64
TextIndexExpansionFactor float64
TieredEvictionEnabled bool
PreferFieldDataWhenIndexHasRawData bool
EnableInterimSegmentIndex bool
TempSegmentIndexFactor float64
ExternalRawDataFactor float64
GrowingMmapEnabled bool
}
func DefaultSegmentLoadingEstimateOptions() SegmentLoadingEstimateOptions {
queryNodeCfg := &paramtable.Get().QueryNodeCfg
return SegmentLoadingEstimateOptions{
MmapVectorField: queryNodeCfg.MmapVectorField.GetAsBool(),
MmapScalarField: queryNodeCfg.MmapScalarField.GetAsBool(),
MmapJSONStats: queryNodeCfg.MmapJSONStats.GetAsBool(),
DeltaDataExpansionFactor: queryNodeCfg.DeltaDataExpansionRate.GetAsFloat(),
JSONKeyStatsExpansionFactor: queryNodeCfg.JSONKeyStatsExpansionFactor.GetAsFloat(),
TextIndexExpansionFactor: queryNodeCfg.TextIndexExpansionFactor.GetAsFloat(),
TieredEvictionEnabled: queryNodeCfg.TieredEvictionEnabled.GetAsBool(),
PreferFieldDataWhenIndexHasRawData: queryNodeCfg.PreferFieldDataWhenIndexHasRawData.GetAsBool(),
EnableInterimSegmentIndex: queryNodeCfg.EnableInterminSegmentIndex.GetAsBool(),
TempSegmentIndexFactor: queryNodeCfg.InterimIndexMemExpandRate.GetAsFloat(),
ExternalRawDataFactor: queryNodeCfg.ExternalCollectionRawDataFactor.GetAsFloat(),
GrowingMmapEnabled: queryNodeCfg.GrowingMmapEnabled.GetAsBool(),
}
}
func estimateTantivyValidityBitmapBytes(numRows int64) uint64 {
if numRows <= 0 {
return 0
}
// TextIndexStats does not carry the null count, so reserve the full
// word-aligned bitmap conservatively. All-valid indexes charge zero actual
// bitmap bytes after loading.
const (
bitsPerWord = uint64(64)
bytesPerWord = uint64(8)
)
words := (uint64(numRows)-1)/bitsPerWord + 1
return words * bytesPerWord
}
// ResolveSegmentEstimateLogs returns the raw and delta metadata consumed by
// the estimators. Internal Storage V3 descriptors are adapted to pathless
// FieldBinlogs; all other SegmentLoadInfo metadata remains unchanged.
func ResolveSegmentEstimateLogs(schema *schemapb.CollectionSchema, loadInfo *querypb.SegmentLoadInfo) ([]*datapb.FieldBinlog, []*datapb.FieldBinlog) {
binlogs := loadInfo.GetBinlogPaths()
deltalogs := loadInfo.GetDeltalogs()
if len(binlogs) > 0 {
return binlogs, deltalogs
}
if loadInfo.GetStorageVersion() != storage.StorageV3 || loadInfo.GetManifestPath() == "" || typeutil.IsExternalCollection(schema) {
return binlogs, deltalogs
}
stats := loadInfo.GetStats()
resource := stats.GetLoadResource()
if resource == nil {
return binlogs, deltalogs
}
groups := resource.GetColumnGroups()
binlogs = make([]*datapb.FieldBinlog, 0, len(groups))
for _, group := range groups {
binlogs = append(binlogs, &datapb.FieldBinlog{
FieldID: group.GetGroupId(),
ChildFields: append([]int64(nil), group.GetFieldIds()...),
Binlogs: []*datapb.Binlog{{
EntriesNum: loadInfo.GetNumOfRows(),
MemorySize: group.GetMemorySize(),
}},
})
}
deltalogs = nil
if deltaSize := stats.GetDeltaBinlogSize(); deltaSize > 0 {
deltalogs = []*datapb.FieldBinlog{{
Binlogs: []*datapb.Binlog{{MemorySize: deltaSize}},
}}
}
return binlogs, deltalogs
}
func EstimateSegmentFinalResource(ctx context.Context, schema *schemapb.CollectionSchema, loadInfo *querypb.SegmentLoadInfo, options SegmentFinalEstimateOptions, runner Runner) (SegmentResourceUsage, error) {
if schema == nil {
return SegmentResourceUsage{}, merr.WrapErrServiceInternalMsg("collection schema is nil")
}
if loadInfo == nil {
return SegmentResourceUsage{}, merr.WrapErrServiceInternalMsg("segment load info is nil")
}
options = options.normalized()
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
if err != nil {
return SegmentResourceUsage{}, err
}
var inevictable SegmentResourceUsage
var evictable SegmentResourceUsage
binlogs, deltalogs := ResolveSegmentEstimateLogs(schema, loadInfo)
id2Binlogs := make(map[int64]*datapb.FieldBinlog, len(binlogs))
for _, fieldBinlog := range binlogs {
id2Binlogs[fieldBinlog.GetFieldID()] = fieldBinlog
}
for _, fieldIndexInfo := range loadInfo.GetIndexInfos() {
if len(fieldIndexInfo.GetIndexFilePaths()) == 0 {
continue
}
fieldID := fieldIndexInfo.GetFieldID()
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
// field might have been dropped, skip its index
mlog.Info(ctx, "skip index for dropped field", mlog.Int64("fieldID", fieldID), mlog.String("name", schema.GetName()))
continue
}
estimate, err := EstimateIndexLoadResourceWithRunner(ctx, fieldSchema, loadInfo, fieldIndexInfo, runner)
if err != nil {
return SegmentResourceUsage{}, errors.Wrapf(err, "failed to estimate final resource usage of index, collection %d, segment %d, indexBuildID %d",
loadInfo.GetCollectionID(),
loadInfo.GetSegmentID(),
fieldIndexInfo.GetBuildID())
}
evictable.MemoryBytes += estimate.FinalMemoryBytes
evictable.DiskBytes += estimate.FinalDiskBytes
if estimate.HasRawData && !options.PreferFieldDataWhenIndexHasRawData {
delete(id2Binlogs, fieldID)
continue
}
if !typeutil.IsVectorType(fieldSchema.GetDataType()) {
continue
}
metricType, err := funcutil.GetAttrByKeyFromRepeatedKV(common.MetricTypeKey, fieldIndexInfo.GetIndexParams())
if err != nil {
return SegmentResourceUsage{}, errors.Wrapf(err, "failed to estimate final resource usage of index, metric type not found, collection %d, segment %d, indexBuildID %d",
loadInfo.GetCollectionID(),
loadInfo.GetSegmentID(),
fieldIndexInfo.GetBuildID())
}
if metricType == metric.BM25 {
delete(id2Binlogs, fieldID)
}
}
for fieldID, fieldBinlog := range id2Binlogs {
fieldIDs := fieldBinlog.GetChildFields()
if len(fieldIDs) == 0 {
fieldIDs = []int64{fieldID}
}
binlogSize := uint64(fieldBinlogMemorySize(fieldBinlog))
var containsTimestampField bool
var doubleMemoryDataField bool
mmapEnabled := true
isVectorType := true
hasLiveField := false
for _, fieldID := range fieldIDs {
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
// field might have been dropped, skip it and continue processing
// other fields in the same column group
mlog.Info(ctx, "skip binlog for dropped field", mlog.Int64("fieldID", fieldID), mlog.String("name", schema.GetName()))
continue
}
hasLiveField = true
isVectorType = isVectorType && typeutil.IsVectorType(fieldSchema.GetDataType())
mmapEnabled = mmapEnabled && isDataMmapEnabled(fieldSchema, options)
containsTimestampField = containsTimestampField || fieldSchema.GetFieldID() == common.TimeStampField
doubleMemoryDataField = doubleMemoryDataField || doubleMemoryDataType(fieldSchema.GetDataType())
}
if !hasLiveField {
continue
}
if containsTimestampField {
timestampSize := int64(0)
for _, binlog := range fieldBinlog.GetBinlogs() {
timestampSize += binlog.GetEntriesNum() * 4
}
inevictable.MemoryBytes += uint64(2 * timestampSize)
}
if isVectorType {
if options.MmapVectorField {
evictable.DiskBytes += binlogSize
} else {
evictable.MemoryBytes += binlogSize
}
continue
}
if !mmapEnabled {
evictable.MemoryBytes += binlogSize
if doubleMemoryDataField {
evictable.MemoryBytes += binlogSize
}
} else {
evictable.DiskBytes += binlogSize
}
}
for _, fieldBinlog := range loadInfo.GetStatslogs() {
inevictable.MemoryBytes += uint64(fieldBinlogMemorySize(fieldBinlog))
}
for _, fieldBinlog := range deltalogs {
expansionFactor := float64(1)
memSize := fieldBinlogMemorySize(fieldBinlog)
if memSize == fieldBinlogDiskSize(fieldBinlog) {
expansionFactor = options.DeltaDataExpansionFactor
}
inevictable.MemoryBytes += uint64(float64(memSize) * expansionFactor)
}
for _, stats := range loadInfo.GetJsonKeyStatsLogs() {
size := uint64(math.Ceil(float64(stats.GetMemorySize()) * options.JSONKeyStatsExpansionFactor))
if options.MmapJSONStats {
evictable.DiskBytes += size
} else {
evictable.MemoryBytes += size
}
}
validityBitmapBytes := estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
for _, stats := range loadInfo.GetTextStatsLogs() {
indexFileBytes := uint64(math.Ceil(float64(stats.GetMemorySize()) * options.TextIndexExpansionFactor))
evictable.MemoryBytes += validityBitmapBytes
if options.MmapScalarField {
evictable.DiskBytes += indexFileBytes
} else {
evictable.MemoryBytes += indexFileBytes
}
}
return SegmentResourceUsage{
MemoryBytes: inevictable.MemoryBytes + applyTieredRatio(evictable.MemoryBytes, options.TieredEvictionEnabled, options.TieredEvictableMemoryCacheRatio),
DiskBytes: inevictable.DiskBytes + applyTieredRatio(evictable.DiskBytes, options.TieredEvictionEnabled, options.TieredEvictableDiskCacheRatio),
}, nil
}
// EstimateSegmentLoadingResource estimates the resource usage of the segment when loading.
// It returns different results depending on whether tiered eviction is enabled:
// - when tiered eviction is enabled, the result is the max resource usage of the segment
// that cannot be managed by caching layer, which should be a subset of the segment inevictable part
// - when tiered eviction is disabled, the result is the max resource usage of both
// the segment evictable and inevictable part
func EstimateSegmentLoadingResource(ctx context.Context, schema *schemapb.CollectionSchema, loadInfo *querypb.SegmentLoadInfo, options SegmentLoadingEstimateOptions, runner Runner) (SegmentResourceUsage, error) {
var segMemoryLoadingSize, segDiskLoadingSize uint64
var indexMemorySize uint64
var fieldGPUMemorySize []uint64
binlogs, deltalogs := ResolveSegmentEstimateLogs(schema, loadInfo)
id2Binlogs := lo.SliceToMap(binlogs, func(fieldBinlog *datapb.FieldBinlog) (int64, *datapb.FieldBinlog) {
return fieldBinlog.GetFieldID(), fieldBinlog
})
schemaHelper, err := typeutil.CreateSchemaHelper(schema)
if err != nil {
mlog.Warn(ctx, "failed to create schema helper", mlog.String("name", schema.GetName()), mlog.Err(err))
return SegmentResourceUsage{}, err
}
indexedFields := make(map[int64]struct{})
// PART 1: calculate size of indexes
for _, fieldIndexInfo := range loadInfo.IndexInfos {
fieldID := fieldIndexInfo.GetFieldID()
if len(fieldIndexInfo.GetIndexFilePaths()) > 0 {
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
// field might have been dropped, skip its index
mlog.Info(ctx, "skip index for dropped field", mlog.Int64("fieldID", fieldID), mlog.String("name", schema.GetName()))
continue
}
indexedFields[fieldID] = struct{}{}
isVectorType := typeutil.IsVectorType(fieldSchema.GetDataType())
estimateResult, err := EstimateIndexLoadResourceWithRunner(ctx, fieldSchema, loadInfo, fieldIndexInfo, runner)
if err != nil {
return SegmentResourceUsage{}, errors.Wrapf(err, "failed to estimate loading resource usage of index, collection %d, segment %d, indexBuildID %d",
loadInfo.GetCollectionID(),
loadInfo.GetSegmentID(),
fieldIndexInfo.GetBuildID())
}
if !options.TieredEvictionEnabled {
indexMemorySize += estimateResult.MaxMemoryBytes
segDiskLoadingSize += estimateResult.MaxDiskBytes
}
if gpuIndexRequiresGpu(fieldIndexInfo.IndexParams) {
fieldGPUMemorySize = append(fieldGPUMemorySize, estimateResult.MaxMemoryBytes)
}
// could skip binlog or
// could be missing for new field or storage v2 group 0
if estimateResult.HasRawData &&
!options.PreferFieldDataWhenIndexHasRawData {
delete(id2Binlogs, fieldID)
continue
}
// BM25 only checks vector datatype
// scalar index does not have metrics type key
if !isVectorType {
continue
}
metricType, err := funcutil.GetAttrByKeyFromRepeatedKV(common.MetricTypeKey, fieldIndexInfo.IndexParams)
if err != nil {
return SegmentResourceUsage{}, errors.Wrapf(err, "failed to estimate loading resource usage of index, metric type not found, collection %d, segment %d, indexBuildID %d",
loadInfo.GetCollectionID(),
loadInfo.GetSegmentID(),
fieldIndexInfo.GetBuildID())
}
// skip raw data for BM25 index
if metricType == metric.BM25 {
delete(id2Binlogs, fieldID)
}
}
}
// PART 2: calculate size of binlogs
for fieldID, fieldBinlog := range id2Binlogs {
fieldIDs := fieldBinlog.GetChildFields()
// legacy default split
if len(fieldIDs) == 0 {
fieldIDs = []int64{fieldID}
}
binlogSize := uint64(loadingFieldBinlogMemorySize(fieldBinlog))
var supportInterimIndexDataType bool
var containsTimestampField bool
var doubleMomoryDataField bool
var legacyNilSchema bool
mmapEnabled := true
isVectorType := true
hasIndex := true
for _, fieldID := range fieldIDs {
// get field schema from fieldID
fieldSchema, err := schemaHelper.GetFieldFromID(fieldID)
if err != nil {
// field might have been dropped, skip it and continue processing
// other fields in the same column group
mlog.Info(ctx, "skip binlog for dropped field", mlog.Int64("fieldID", fieldID), mlog.String("name", schema.GetName()))
continue
}
if _, ok := indexedFields[fieldID]; !ok {
hasIndex = false
}
// missing mapping, shall be "0" group for storage v2
if fieldSchema == nil {
if !options.TieredEvictionEnabled {
segMemoryLoadingSize += binlogSize
}
legacyNilSchema = true
break
}
supportInterimIndexDataType = supportInterimIndexDataType || supportInterimIndexDataTypeForLoading(fieldSchema.GetDataType())
isVectorType = isVectorType && typeutil.IsVectorType(fieldSchema.GetDataType())
mmapEnabled = mmapEnabled && isLoadingDataMmapEnabled(fieldSchema, options)
containsTimestampField = containsTimestampField || fieldSchema.GetFieldID() == common.TimeStampField
doubleMomoryDataField = doubleMomoryDataField || doubleMemoryDataType(fieldSchema.GetDataType())
}
// legacy v2 segment without children
if legacyNilSchema {
continue
}
if !hasIndex {
if !options.TieredEvictionEnabled {
interimIndexEnable := options.EnableInterimSegmentIndex && !options.GrowingMmapEnabled && supportInterimIndexDataType
if interimIndexEnable {
segMemoryLoadingSize += uint64(float64(binlogSize) * options.TempSegmentIndexFactor)
}
}
}
if isVectorType {
if options.MmapVectorField {
if !options.TieredEvictionEnabled {
segDiskLoadingSize += binlogSize
}
} else {
if !options.TieredEvictionEnabled {
segMemoryLoadingSize += binlogSize
}
}
continue
}
// timestamp field double in InsertRecord & TimestampIndex
if containsTimestampField {
timestampSize := lo.SumBy(fieldBinlog.GetBinlogs(), func(binlog *datapb.Binlog) int64 {
return binlog.GetEntriesNum() * 4
})
segMemoryLoadingSize += 2 * uint64(timestampSize)
}
if !mmapEnabled {
if !options.TieredEvictionEnabled {
segMemoryLoadingSize += binlogSize
if doubleMomoryDataField {
segMemoryLoadingSize += binlogSize
}
}
} else {
if !options.TieredEvictionEnabled {
segDiskLoadingSize += uint64(loadingFieldBinlogMemorySize(fieldBinlog))
}
}
}
// PART 2.5: external segment adjustments
//
// External segments carry pre-computed MemorySize in fake binlogs (from
// DataNode Take sampling). Adjust the memory estimate for two external-
// specific behaviors:
// 1. Non-lazy path without tiered eviction: apply ExternalRawDataFactor
// to cover the peak transient memory during download + decompress +
// Arrow deserialize (normal packed segments do not have this peak
// because their binlogs are already in Arrow IPC format). When tiered
// eviction is enabled, the caching layer reserves this loading
// overhead, so Go must not reserve the same raw-data margin again.
// 2. Full-lazy path (all external fields warmup=disable): no eager
// load, so subtract the raw data size that PART 2 added.
// Also propagate EstimatedBytesPerRow to the C++ ManifestGroupTranslator
// so the tiered-cache layer sizes chunks correctly.
if typeutil.IsExternalCollection(schema) && loadInfo.GetNumOfRows() > 0 {
var fakeBinlogMemSize int64
for _, fb := range binlogs {
fakeBinlogMemSize += loadingFieldBinlogMemorySize(fb)
}
loadInfo.EstimatedBytesPerRow = fakeBinlogMemSize / loadInfo.GetNumOfRows()
if isExternalCollectionLazyLoadForLoading(schema) {
// Full-lazy -> zero eager load. Undo PART 2's rawSize addition.
// Safety factor does not apply: no peak to cover.
if segMemoryLoadingSize >= uint64(fakeBinlogMemSize) {
segMemoryLoadingSize -= uint64(fakeBinlogMemSize)
} else {
segMemoryLoadingSize = 0
}
} else if !options.TieredEvictionEnabled {
if factor := options.ExternalRawDataFactor; factor > 1.0 {
// Non-lazy -> add peak margin on top of rawSize that PART 2 added.
segMemoryLoadingSize += uint64(float64(fakeBinlogMemSize) * (factor - 1.0))
}
}
}
// PART 3: calculate size of stats data
// stats data isn't managed by the caching layer, so its size should always be included,
// regardless of the tiered eviction value
for _, fieldBinlog := range loadInfo.GetStatslogs() {
segMemoryLoadingSize += uint64(loadingFieldBinlogMemorySize(fieldBinlog))
}
// PART 4: calculate size of delete data
// delete data isn't managed by the caching layer, so its size should always be included,
// regardless of the tiered eviction value
for _, fieldBinlog := range deltalogs {
// MemorySize of filedBinlog is the actual size in memory, but we should also consider
// the memcpy from golang to cpp side, so the expansionFactor is set to 2.
expansionFactor := float64(2)
memSize := loadingFieldBinlogMemorySize(fieldBinlog)
// Note: If MemorySize == DiskSize, it means the segment comes from Milvus 2.3,
// MemorySize is actually compressed DiskSize of deltalog, so we'll fallback to use
// deltaExpansionFactor to compromise the compression ratio.
if memSize == loadingFieldBinlogDiskSize(fieldBinlog) {
expansionFactor = options.DeltaDataExpansionFactor
}
segMemoryLoadingSize += uint64(float64(memSize) * expansionFactor)
}
// PART 5: calculate size of json key stats data
for _, jsonKeyStats := range loadInfo.GetJsonKeyStatsLogs() {
if options.MmapJSONStats {
if !options.TieredEvictionEnabled {
segDiskLoadingSize += uint64(float64(jsonKeyStats.GetMemorySize()) * options.JSONKeyStatsExpansionFactor)
}
} else {
if !options.TieredEvictionEnabled {
segMemoryLoadingSize += uint64(float64(jsonKeyStats.GetMemorySize()) * options.JSONKeyStatsExpansionFactor)
}
}
}
// per struct memory size, used to keep mapping between row id and element id
var structArrayOffsetsSize uint64
// PART 6: calculate size of struct array offsets
// The memory size is 4 * row_count + 4 * total_element_count
// We cannot easily get the element count, so we estimate it by the row count * 10
rowCount := uint64(loadInfo.GetNumOfRows())
for range len(schema.GetStructArrayFields()) {
structArrayOffsetsSize += 4*rowCount + 4*rowCount*10
}
// PART 7: calculate size of text index stats data
// text index data is managed by the caching layer when tiered eviction is enabled,
// so it only needs to be included when tiered eviction is disabled.
// Text match index mmap is driven by scalar_field_enable_mmap (same as raw scalar data).
// memory_size is the sum of uploaded Tantivy index files, including sparse null sidecars.
// The materialized word-aligned validity bitmap is separate heap memory, and
// TextIndexExpansionFactor applies only to the index file bytes.
validityBitmapBytes := estimateTantivyValidityBitmapBytes(loadInfo.GetNumOfRows())
for _, textStats := range loadInfo.GetTextStatsLogs() {
if options.TieredEvictionEnabled {
continue
}
indexFileBytes := uint64(float64(textStats.GetMemorySize()) * options.TextIndexExpansionFactor)
segMemoryLoadingSize += validityBitmapBytes
if options.MmapScalarField {
segDiskLoadingSize += indexFileBytes
} else {
segMemoryLoadingSize += indexFileBytes
}
}
return SegmentResourceUsage{
MemoryBytes: segMemoryLoadingSize + indexMemorySize + structArrayOffsetsSize,
DiskBytes: segDiskLoadingSize,
FieldGPUMemoryBytes: fieldGPUMemorySize,
}, nil
}
func (options SegmentFinalEstimateOptions) normalized() SegmentFinalEstimateOptions {
if options.DeltaDataExpansionFactor <= 0 {
options.DeltaDataExpansionFactor = 1
}
if options.JSONKeyStatsExpansionFactor <= 0 {
options.JSONKeyStatsExpansionFactor = 1
}
if options.TextIndexExpansionFactor <= 0 {
options.TextIndexExpansionFactor = 1
}
if options.TieredEvictableMemoryCacheRatio < 0 || options.TieredEvictableMemoryCacheRatio > 1 {
options.TieredEvictableMemoryCacheRatio = 1
}
if options.TieredEvictableDiskCacheRatio < 0 || options.TieredEvictableDiskCacheRatio > 1 {
options.TieredEvictableDiskCacheRatio = 1
}
return options
}
func isDataMmapEnabled(field *schemapb.FieldSchema, options SegmentFinalEstimateOptions) bool {
if enabled, ok := common.IsMmapDataEnabled(field.GetTypeParams()...); ok {
return enabled
}
if typeutil.IsVectorType(field.GetDataType()) {
return options.MmapVectorField
}
return options.MmapScalarField
}
func isLoadingDataMmapEnabled(field *schemapb.FieldSchema, options SegmentLoadingEstimateOptions) bool {
enableMmap, exist := common.IsMmapDataEnabled(field.GetTypeParams()...)
if exist {
return enableMmap
}
if typeutil.IsVectorType(field.GetDataType()) {
return options.MmapVectorField
}
return options.MmapScalarField
}
func doubleMemoryDataType(dataType schemapb.DataType) bool {
return dataType == schemapb.DataType_String ||
dataType == schemapb.DataType_VarChar ||
dataType == schemapb.DataType_JSON
}
func supportInterimIndexDataTypeForLoading(dataType schemapb.DataType) bool {
return dataType == schemapb.DataType_FloatVector ||
dataType == schemapb.DataType_SparseFloatVector ||
dataType == schemapb.DataType_Float16Vector ||
dataType == schemapb.DataType_BFloat16Vector
}
func isExternalCollectionLazyLoadForLoading(schema *schemapb.CollectionSchema) bool {
resolver := typeutil.NewStorageColumnResolver(schema)
for _, field := range schema.GetFields() {
if !resolver.IsSourceDataField(field) {
continue
}
if resolver.IsMilvusTable() && field.GetIsPrimaryKey() {
// Real-PK milvus-table segments always load the source PK column and
// source insert timestamps eagerly so source deltas preserve
// delete/reinsert ordering.
return false
}
policy := getFieldWarmupPolicyForLoading(field)
if policy != common.WarmupDisable {
return false
}
}
return true
}
func getFieldWarmupPolicyForLoading(field *schemapb.FieldSchema) string {
policy, exist := common.GetWarmupPolicy(field.GetTypeParams()...)
if exist {
return policy
}
if typeutil.IsVectorType(field.GetDataType()) {
return paramtable.Get().QueryNodeCfg.TieredWarmupVectorField.GetValue()
}
return paramtable.Get().QueryNodeCfg.TieredWarmupScalarField.GetValue()
}
func gpuIndexRequiresGpu(indexParams []*commonpb.KeyValuePair) bool {
indexParamMap := funcutil.KeyValuePair2Map(indexParams)
indexType := indexParamMap[common.IndexTypeKey]
switch indexType {
case "GPU_CAGRA", "GPU_CUVS_CAGRA":
case "GPU_BRUTE_FORCE", "GPU_CUVS_BRUTE_FORCE",
"GPU_IVF_FLAT", "GPU_CUVS_IVF_FLAT",
"GPU_IVF_PQ", "GPU_CUVS_IVF_PQ":
return true
default:
return false
}
err := indexparams.AppendPrepareLoadParams(paramtable.Get(), indexParamMap)
if err != nil {
mlog.Warn(context.TODO(), "failed to append prepare load params for gpu index resource check",
mlog.String("indexType", indexType),
mlog.Err(err))
}
adaptForCPU, ok := indexParamMap["adapt_for_cpu"]
if ok {
enabled, err := strconv.ParseBool(adaptForCPU)
if err == nil && enabled {
return false
}
}
return true
}
func fieldBinlogMemorySize(fieldBinlog *datapb.FieldBinlog) int64 {
fieldSize := int64(0)
for _, binlog := range fieldBinlog.GetBinlogs() {
if binlog.GetMemorySize() > 0 {
fieldSize += binlog.GetMemorySize()
continue
}
fieldSize += binlog.GetLogSize()
}
return fieldSize
}
func loadingFieldBinlogMemorySize(fieldBinlog *datapb.FieldBinlog) int64 {
fieldSize := int64(0)
for _, binlog := range fieldBinlog.Binlogs {
fieldSize += binlog.GetMemorySize()
}
return fieldSize
}
func fieldBinlogDiskSize(fieldBinlog *datapb.FieldBinlog) int64 {
fieldSize := int64(0)
for _, binlog := range fieldBinlog.GetBinlogs() {
fieldSize += binlog.GetLogSize()
}
return fieldSize
}
func loadingFieldBinlogDiskSize(fieldBinlog *datapb.FieldBinlog) int64 {
fieldSize := int64(0)
for _, binlog := range fieldBinlog.Binlogs {
fieldSize += binlog.GetLogSize()
}
return fieldSize
}
func applyTieredRatio(value uint64, enabled bool, ratio float64) uint64 {
if !enabled {
return value
}
return uint64(math.Ceil(float64(value) * ratio))
}
internal/util/segcore/loadresource/segment_test.go
+402
View File
@@ -0,0 +1,402 @@
// 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 loadresource
import (
"context"
"testing"
"github.com/stretchr/testify/require"
"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/datapb"
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
"github.com/milvus-io/milvus/pkg/v3/util/metric"
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
)
func TestEstimateTantivyValidityBitmapBytesWordAligned(t *testing.T) {
require.EqualValues(t, 0, estimateTantivyValidityBitmapBytes(-1))
require.EqualValues(t, 0, estimateTantivyValidityBitmapBytes(0))
require.EqualValues(t, 8, estimateTantivyValidityBitmapBytes(1))
require.EqualValues(t, 8, estimateTantivyValidityBitmapBytes(64))
require.EqualValues(t, 16, estimateTantivyValidityBitmapBytes(65))
require.EqualValues(t, 24, estimateTantivyValidityBitmapBytes(129))
}
func TestEstimateSegmentFinalResourceUsesBinlogStatsAndTieredRatio(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_final_estimate",
Fields: []*schemapb.FieldSchema{
{FieldID: common.TimeStampField, Name: common.TimeStampFieldName, DataType: schemapb.DataType_Int64},
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}},
{FieldID: 102, Name: "json", DataType: schemapb.DataType_JSON},
{FieldID: 103, Name: "bm25", DataType: schemapb.DataType_SparseFloatVector},
},
}
loadInfo := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 10,
BinlogPaths: []*datapb.FieldBinlog{
{
FieldID: common.TimeStampField,
Binlogs: []*datapb.Binlog{
{MemorySize: 80, LogSize: 50, EntriesNum: 10},
},
},
{
FieldID: 101,
Binlogs: []*datapb.Binlog{
{LogSize: 160},
},
},
{
FieldID: 102,
Binlogs: []*datapb.Binlog{
{MemorySize: 200, LogSize: 120},
},
},
},
Statslogs: []*datapb.FieldBinlog{
{Binlogs: []*datapb.Binlog{{MemorySize: 30, LogSize: 20}}},
},
Deltalogs: []*datapb.FieldBinlog{
{Binlogs: []*datapb.Binlog{{MemorySize: 40, LogSize: 40}}},
},
Bm25Logs: []*datapb.FieldBinlog{
{FieldID: 103, Binlogs: []*datapb.Binlog{{MemorySize: 40, LogSize: 20}}},
},
JsonKeyStatsLogs: map[int64]*datapb.JsonKeyStats{
102: {MemorySize: 100, LogSize: 90},
},
TextStatsLogs: map[int64]*datapb.TextIndexStats{
102: {MemorySize: 50, LogSize: 40},
},
}
for _, test := range []struct {
name string
tieredEnabled bool
expectedMemory uint64
expectedDisk uint64
}{
{name: "disabled uses full evictable resource", expectedMemory: 778, expectedDisk: 310},
{name: "enabled applies cache ratios", tieredEnabled: true, expectedMemory: 484, expectedDisk: 78},
} {
t.Run(test.name, func(t *testing.T) {
usage, err := EstimateSegmentFinalResource(context.Background(), schema, loadInfo, SegmentFinalEstimateOptions{
MmapVectorField: true,
MmapScalarField: false,
MmapJSONStats: true,
TieredEvictionEnabled: test.tieredEnabled,
TieredEvictableMemoryCacheRatio: 0.5,
TieredEvictableDiskCacheRatio: 0.25,
DeltaDataExpansionFactor: 2,
JSONKeyStatsExpansionFactor: 1.5,
TextIndexExpansionFactor: 2,
PreferFieldDataWhenIndexHasRawData: false,
}, nil)
require.NoError(t, err)
require.Equal(t, test.expectedMemory, usage.MemoryBytes)
require.Equal(t, test.expectedDisk, usage.DiskBytes)
})
}
}
func TestEstimateSegmentLoadingResourceAccountsForTantivyValidityBitmap(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_text_loading_estimate",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "text", DataType: schemapb.DataType_VarChar},
},
}
loadInfo := &querypb.SegmentLoadInfo{
NumOfRows: 65,
TextStatsLogs: map[int64]*datapb.TextIndexStats{
101: {FieldID: 101, MemorySize: 100},
},
}
for _, test := range []struct {
name string
mmap bool
tiered bool
expectedMemory uint64
expectedDisk uint64
}{
{name: "non-mmap", expectedMemory: 116},
{name: "mmap", mmap: true, expectedMemory: 16, expectedDisk: 100},
{name: "tiered skips loading reservation", tiered: true},
} {
t.Run(test.name, func(t *testing.T) {
usage, err := EstimateSegmentLoadingResource(context.Background(), schema, loadInfo, SegmentLoadingEstimateOptions{
MmapScalarField: test.mmap,
TieredEvictionEnabled: test.tiered,
TextIndexExpansionFactor: 1,
}, nil)
require.NoError(t, err)
require.Equal(t, test.expectedMemory, usage.MemoryBytes)
require.Equal(t, test.expectedDisk, usage.DiskBytes)
})
}
}
func TestEstimateSegmentFinalResourceAllowsZeroTieredRatio(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_zero_tiered_ratio",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}},
},
}
loadInfo := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 10,
BinlogPaths: []*datapb.FieldBinlog{
{
FieldID: 100,
Binlogs: []*datapb.Binlog{
{MemorySize: 100, LogSize: 80},
},
},
{
FieldID: 101,
Binlogs: []*datapb.Binlog{
{MemorySize: 200, LogSize: 160},
},
},
},
Statslogs: []*datapb.FieldBinlog{
{Binlogs: []*datapb.Binlog{{MemorySize: 30, LogSize: 20}}},
},
}
usage, err := EstimateSegmentFinalResource(context.Background(), schema, loadInfo, SegmentFinalEstimateOptions{
MmapVectorField: true,
TieredEvictionEnabled: true,
TieredEvictableMemoryCacheRatio: 0,
TieredEvictableDiskCacheRatio: 0,
}, nil)
require.NoError(t, err)
require.EqualValues(t, 30, usage.MemoryBytes)
require.Zero(t, usage.DiskBytes)
}
func TestEstimateSegmentFinalResourceOnlyExpandsLegacyDelta(t *testing.T) {
schema := &schemapb.CollectionSchema{Name: "test_delta_estimate"}
for _, test := range []struct {
name string
memorySize int64
logSize int64
expectedMemory uint64
}{
{
name: "storage v3 summary has no disk size",
memorySize: 100,
expectedMemory: 100,
},
{
name: "legacy compressed size uses expansion",
memorySize: 100,
logSize: 100,
expectedMemory: 300,
},
} {
t.Run(test.name, func(t *testing.T) {
usage, err := EstimateSegmentFinalResource(context.Background(), schema, &querypb.SegmentLoadInfo{
Deltalogs: []*datapb.FieldBinlog{{
Binlogs: []*datapb.Binlog{{MemorySize: test.memorySize, LogSize: test.logSize}},
}},
}, SegmentFinalEstimateOptions{DeltaDataExpansionFactor: 3}, nil)
require.NoError(t, err)
require.Equal(t, test.expectedMemory, usage.MemoryBytes)
require.Zero(t, usage.DiskBytes)
})
}
}
func TestEstimateSegmentFinalResourceSkipsDroppedFields(t *testing.T) {
schema := &schemapb.CollectionSchema{
Name: "test_dropped_fields",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
},
}
t.Run("index", func(t *testing.T) {
runnerCalled := false
usage, err := EstimateSegmentFinalResource(context.Background(), schema, &querypb.SegmentLoadInfo{
IndexInfos: []*querypb.FieldIndexInfo{
{FieldID: 999, IndexFilePaths: []string{"index"}},
},
BinlogPaths: []*datapb.FieldBinlog{
{FieldID: 100, Binlogs: []*datapb.Binlog{{MemorySize: 64}}},
},
}, SegmentFinalEstimateOptions{}, func(fn func() error) error {
runnerCalled = true
return fn()
})
require.NoError(t, err)
require.False(t, runnerCalled)
require.EqualValues(t, 64, usage.MemoryBytes)
require.Zero(t, usage.DiskBytes)
})
t.Run("binlog group with only dropped fields", func(t *testing.T) {
usage, err := EstimateSegmentFinalResource(context.Background(), schema, &querypb.SegmentLoadInfo{
BinlogPaths: []*datapb.FieldBinlog{
{
FieldID: 200,
ChildFields: []int64{999},
Binlogs: []*datapb.Binlog{{MemorySize: 128}},
},
{
FieldID: 998,
Binlogs: []*datapb.Binlog{{MemorySize: 128}},
},
},
}, SegmentFinalEstimateOptions{}, nil)
require.NoError(t, err)
require.Zero(t, usage.MemoryBytes)
require.Zero(t, usage.DiskBytes)
})
t.Run("binlog group with live and dropped fields", func(t *testing.T) {
usage, err := EstimateSegmentFinalResource(context.Background(), schema, &querypb.SegmentLoadInfo{
BinlogPaths: []*datapb.FieldBinlog{
{
FieldID: 200,
ChildFields: []int64{999, 100},
Binlogs: []*datapb.Binlog{{MemorySize: 128}},
},
},
}, SegmentFinalEstimateOptions{MmapVectorField: true}, nil)
require.NoError(t, err)
require.EqualValues(t, 128, usage.MemoryBytes)
require.Zero(t, usage.DiskBytes)
})
}
func TestEstimateSegmentLoadingResourceSkipsCAGRAGPUMemoryWhenAdaptForCPU(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_gpu_adapt_for_cpu",
Fields: []*schemapb.FieldSchema{
{FieldID: 100, Name: "pk", DataType: schemapb.DataType_Int64, IsPrimaryKey: true},
{FieldID: 101, Name: "vec", DataType: schemapb.DataType_FloatVector, TypeParams: []*commonpb.KeyValuePair{{Key: common.DimKey, Value: "4"}}},
},
}
loadInfo := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 10,
IndexInfos: []*querypb.FieldIndexInfo{
{
FieldID: 101,
IndexID: 10,
BuildID: 11,
IndexFilePaths: []string{"index"},
IndexParams: []*commonpb.KeyValuePair{
{Key: common.IndexTypeKey, Value: "GPU_CAGRA"},
{Key: common.MetricTypeKey, Value: metric.L2},
{Key: "adapt_for_cpu", Value: "true"},
},
},
},
}
usage, err := EstimateSegmentLoadingResource(context.Background(), schema, loadInfo, SegmentLoadingEstimateOptions{}, func(func() error) error {
return nil
})
require.NoError(t, err)
require.Empty(t, usage.FieldGPUMemoryBytes)
}
func TestEstimateSegmentLoadingResourceKeepsMilvusTableRealPKEager(t *testing.T) {
paramtable.Init()
schema := &schemapb.CollectionSchema{
Name: "test_milvus_table_real_pk",
ExternalSource: "s3://bucket/source",
ExternalSpec: `{"format":"milvus-table"}`,
Fields: []*schemapb.FieldSchema{
{
FieldID: 100,
Name: "pk",
DataType: schemapb.DataType_Int64,
IsPrimaryKey: true,
ExternalField: "pk",
TypeParams: []*commonpb.KeyValuePair{
{Key: common.WarmupKey, Value: common.WarmupDisable},
},
},
{
FieldID: 101,
Name: "vec",
DataType: schemapb.DataType_FloatVector,
ExternalField: "vec",
TypeParams: []*commonpb.KeyValuePair{
{Key: common.DimKey, Value: "4"},
{Key: common.WarmupKey, Value: common.WarmupDisable},
},
},
},
}
loadInfo := &querypb.SegmentLoadInfo{
CollectionID: 1,
PartitionID: 2,
SegmentID: 3,
NumOfRows: 10,
BinlogPaths: []*datapb.FieldBinlog{
{
FieldID: 0,
ChildFields: []int64{100, 101},
Binlogs: []*datapb.Binlog{
{MemorySize: 1000, LogSize: 1000, EntriesNum: 10},
},
},
},
}
usage, err := EstimateSegmentLoadingResource(context.Background(), schema, loadInfo, SegmentLoadingEstimateOptions{
ExternalRawDataFactor: 2,
}, nil)
require.NoError(t, err)
require.EqualValues(t, 2000, usage.MemoryBytes)
require.Zero(t, usage.DiskBytes)
require.EqualValues(t, 100, loadInfo.GetEstimatedBytesPerRow())
}
pkg/metrics/metrics_test.go
+28
View File
@@ -25,9 +25,12 @@ import (
"testing"
"github.com/prometheus/client_golang/prometheus"
"github.com/prometheus/client_golang/prometheus/testutil"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/sync/errgroup"
"github.com/milvus-io/milvus/pkg/v3/common"
)
func TestRegisterMetrics(t *testing.T) {
@@ -48,6 +51,31 @@ func TestRegisterMetrics(t *testing.T) {
})
}
func TestRegisterQueryCoordInitializesLoadDemandMetrics(t *testing.T) {
QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(common.DefaultResourceGroupName)
QueryCoordLoadDemandDiskBytes.DeleteLabelValues(common.DefaultResourceGroupName)
t.Cleanup(func() {
QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(common.DefaultResourceGroupName)
QueryCoordLoadDemandDiskBytes.DeleteLabelValues(common.DefaultResourceGroupName)
})
registry := prometheus.NewRegistry()
RegisterQueryCoord(registry)
require.NoError(t, testutil.GatherAndCompare(
registry,
strings.NewReader(`# HELP milvus_querycoord_load_demand_disk_bytes cumulative estimated disk bytes required by load configuration changes successfully broadcast after resource precheck
# TYPE milvus_querycoord_load_demand_disk_bytes counter
milvus_querycoord_load_demand_disk_bytes{rg="__default_resource_group"} 0
# HELP milvus_querycoord_load_demand_memory_bytes cumulative estimated memory bytes required by load configuration changes successfully broadcast after resource precheck
# TYPE milvus_querycoord_load_demand_memory_bytes counter
milvus_querycoord_load_demand_memory_bytes{rg="__default_resource_group"} 0
`),
"milvus_querycoord_load_demand_memory_bytes",
"milvus_querycoord_load_demand_disk_bytes",
))
}
func TestGetRegisterer(t *testing.T) {
register := GetRegisterer()
assert.NotNil(t, register)
pkg/metrics/querycoord_metrics.go
+22
View File
@@ -21,6 +21,7 @@ import (
"github.com/prometheus/client_golang/prometheus"
"github.com/milvus-io/milvus/pkg/v3/common"
"github.com/milvus-io/milvus/pkg/v3/util/typeutil"
)
@@ -106,6 +107,22 @@ var (
Help: "the number of tasks in QueryCoord's scheduler",
}, []string{QueryCoordTaskType})
QueryCoordLoadDemandMemoryBytes = prometheus.NewCounterVec(
prometheus.CounterOpts{
Namespace: milvusNamespace,
Subsystem: typeutil.QueryCoordRole,
Name: "load_demand_memory_bytes",
Help: "cumulative estimated memory bytes required by load configuration changes successfully broadcast after resource precheck",
}, []string{ResourceGroupLabelName})
QueryCoordLoadDemandDiskBytes = prometheus.NewCounterVec(
prometheus.CounterOpts{
Namespace: milvusNamespace,
Subsystem: typeutil.QueryCoordRole,
Name: "load_demand_disk_bytes",
Help: "cumulative estimated disk bytes required by load configuration changes successfully broadcast after resource precheck",
}, []string{ResourceGroupLabelName})
QueryCoordNumQueryNodes = prometheus.NewGaugeVec(
prometheus.GaugeOpts{
Namespace: milvusNamespace,
@@ -180,6 +197,9 @@ var (
// RegisterQueryCoord registers QueryCoord metrics
func RegisterQueryCoord(registry *prometheus.Registry) {
QueryCoordLoadDemandMemoryBytes.WithLabelValues(common.DefaultResourceGroupName).Add(0)
QueryCoordLoadDemandDiskBytes.WithLabelValues(common.DefaultResourceGroupName).Add(0)
registry.MustRegister(QueryCoordNumCollections)
registry.MustRegister(QueryCoordNumPartitions)
registry.MustRegister(QueryCoordLoadCount)
@@ -187,6 +207,8 @@ func RegisterQueryCoord(registry *prometheus.Registry) {
registry.MustRegister(QueryCoordLoadLatency)
registry.MustRegister(QueryCoordReleaseLatency)
registry.MustRegister(QueryCoordTaskNum)
registry.MustRegister(QueryCoordLoadDemandMemoryBytes)
registry.MustRegister(QueryCoordLoadDemandDiskBytes)
registry.MustRegister(QueryCoordNumQueryNodes)
registry.MustRegister(QueryCoordCurrentTargetCheckpointUnixSeconds)
registry.MustRegister(QueryCoordCurrentTargetAllReplicasCheckpointUnixSeconds)
pkg/util/metricsinfo/metrics_info.go
+10 -3
View File
@@ -276,12 +276,19 @@ type QueryNodeCollectionMetrics struct {
CollectionRows map[int64]int64
}
// QueryNodeLocalStorageMetrics describes the QueryNode local-storage quota and directory usage.
type QueryNodeLocalStorageMetrics struct {
CapacityBytes int64 `json:"capacity_bytes"`
UsedBytes int64 `json:"used_bytes"`
}
// QueryNodeInfos implements ComponentInfos
type QueryNodeInfos struct {
BaseComponentInfos
SystemConfigurations QueryNodeConfiguration `json:"system_configurations"`
QuotaMetrics *QueryNodeQuotaMetrics `json:"quota_metrics"`
CollectionMetrics *QueryNodeCollectionMetrics `json:"collection_metrics"`
SystemConfigurations QueryNodeConfiguration `json:"system_configurations"`
QuotaMetrics *QueryNodeQuotaMetrics `json:"quota_metrics"`
CollectionMetrics *QueryNodeCollectionMetrics `json:"collection_metrics"`
LocalStorage *QueryNodeLocalStorageMetrics `json:"local_storage,omitempty"`
}
// QueryCoordConfiguration records the configuration of QueryCoord.
pkg/util/metricsinfo/metrics_info_test.go
+4
View File
@@ -88,6 +88,10 @@ func TestQueryNodeInfos_Codec(t *testing.T) {
SystemConfigurations: QueryNodeConfiguration{
SimdType: "avx2",
},
LocalStorage: &QueryNodeLocalStorageMetrics{
CapacityBytes: 1600,
UsedBytes: 2,
},
}
s, err := MarshalComponentInfos(infos1)
assert.Equal(t, nil, err)
pkg/util/paramtable/component_param.go
+41
View File
@@ -3555,6 +3555,11 @@ type queryCoordConfig struct {
ResourceExhaustionPenaltyDuration ParamItem `refreshable:"true"`
ResourceExhaustionCleanupInterval ParamItem `refreshable:"true"`
AutoscalePrecheckEnabled ParamItem `refreshable:"true"`
AutoscaleEnabled ParamItem `refreshable:"true"`
AutoscaleMaxMemoryLimit ParamItem `refreshable:"true"`
AutoscaleMaxDiskLimit ParamItem `refreshable:"true"`
UpdateTargetNeedSegmentDataReady ParamItem `refreshable:"true"`
AutoWarmupForNonPKIsolationCollection ParamItem `refreshable:"false"`
@@ -4264,6 +4269,42 @@ Set to 0 to disable the penalty period.`,
}
p.ResourceExhaustionCleanupInterval.Init(base.mgr)
p.AutoscalePrecheckEnabled = ParamItem{
Key: "queryCoord.autoscale.precheckEnabled",
Version: "3.0.0",
DefaultValue: "false",
Doc: "Whether QueryCoord performs load resource precheck before committing load configuration changes and publishes demand counters after successful broadcasts. With tiered eviction enabled, only demand estimation is performed; capacity checks are skipped.",
Export: true,
}
p.AutoscalePrecheckEnabled.Init(base.mgr)
p.AutoscaleEnabled = ParamItem{
Key: "queryCoord.autoscale.enabled",
Version: "3.0.0",
DefaultValue: "false",
Doc: "Whether QueryCoord load resource precheck can use autoscale upper bounds to admit a load request when current QueryNode capacity is insufficient. This is effective only when precheckEnabled is true and tiered eviction is disabled.",
Export: true,
}
p.AutoscaleEnabled.Init(base.mgr)
p.AutoscaleMaxMemoryLimit = ParamItem{
Key: "queryCoord.autoscale.maxMemoryLimit",
Version: "3.0.0",
DefaultValue: "0",
Doc: "Maximum global QueryNode memory capacity after autoscale, unit: GiB (1 GiB = 1024^3 bytes). 0 means autoscale cannot add memory capacity.",
Export: true,
}
p.AutoscaleMaxMemoryLimit.Init(base.mgr)
p.AutoscaleMaxDiskLimit = ParamItem{
Key: "queryCoord.autoscale.maxDiskLimit",
Version: "3.0.0",
DefaultValue: "0",
Doc: "Maximum sum of Worker QueryNode LOCAL_STORAGE_SIZE after autoscale, unit: GiB (1 GiB = 1024^3 bytes). 0 means autoscale cannot add disk capacity.",
Export: true,
}
p.AutoscaleMaxDiskLimit.Init(base.mgr)
p.UpdateTargetNeedSegmentDataReady = ParamItem{
Key: "queryCoord.updateTargetNeedSegmentDataReady",
Version: "2.6.8",
pkg/util/paramtable/component_param_test.go
+13
View File
@@ -973,6 +973,19 @@ func TestComponentParam(t *testing.T) {
assert.Equal(t, 100, Params.BalanceCheckCollectionMaxCount.GetAsInt())
assert.Equal(t, 30, Params.ResourceExhaustionPenaltyDuration.GetAsInt())
assert.Equal(t, 10, Params.ResourceExhaustionCleanupInterval.GetAsInt())
assert.False(t, Params.AutoscalePrecheckEnabled.GetAsBool())
assert.False(t, Params.AutoscaleEnabled.GetAsBool())
assert.Equal(t, int64(0), Params.AutoscaleMaxMemoryLimit.GetAsInt64())
assert.Equal(t, int64(0), Params.AutoscaleMaxDiskLimit.GetAsInt64())
params.Save("queryCoord.autoscale.precheckEnabled", "true")
params.Save("queryCoord.autoscale.enabled", "true")
params.Save("queryCoord.autoscale.maxMemoryLimit", "1024")
params.Save("queryCoord.autoscale.maxDiskLimit", "2048")
assert.True(t, Params.AutoscalePrecheckEnabled.GetAsBool())
assert.True(t, Params.AutoscaleEnabled.GetAsBool())
assert.Equal(t, int64(1024), Params.AutoscaleMaxMemoryLimit.GetAsInt64())
assert.Equal(t, int64(2048), Params.AutoscaleMaxDiskLimit.GetAsInt64())
})
t.Run("test queryNodeConfig", func(t *testing.T) {