diff --git a/configs/milvus.yaml b/configs/milvus.yaml index d0a2da8a5e..1c963034d5 100644 --- a/configs/milvus.yaml +++ b/configs/milvus.yaml @@ -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 diff --git a/deployments/monitor/grafana/README.md b/deployments/monitor/grafana/README.md index 8e36958c2e..e846f712b8 100644 --- a/deployments/monitor/grafana/README.md +++ b/deployments/monitor/grafana/README.md @@ -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.
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:
``` 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]))) ```
avg:
``` 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:
``` sum(milvus_querycoord_load_demand_memory_bytes{app_kubernetes_io_instance=~"$instance", app_kubernetes_io_name="$app_name", namespace="$namespace"}) by (rg) / 1024 / 1024 ```
disk:
``` 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. | diff --git a/deployments/monitor/grafana/milvus-dashboard.json b/deployments/monitor/grafana/milvus-dashboard.json index 34a33c1899..f3f724e4b4 100644 --- a/deployments/monitor/grafana/milvus-dashboard.json +++ b/deployments/monitor/grafana/milvus-dashboard.json @@ -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": "" -} \ No newline at end of file +} diff --git a/docs/design-docs/design_docs/20260721-querycoord-autoscale-load-resource-precheck.md b/docs/design-docs/design_docs/20260721-querycoord-autoscale-load-resource-precheck.md new file mode 100644 index 0000000000..3c1026d9f1 --- /dev/null +++ b/docs/design-docs/design_docs/20260721-querycoord-autoscale-load-resource-precheck.md @@ -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"] = "" +``` + +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. diff --git a/internal/core/src/segcore/load_index_c.cpp b/internal/core/src/segcore/load_index_c.cpp index 6132a2998d..7ea4d24c79 100644 --- a/internal/core/src/segcore/load_index_c.cpp +++ b/internal/core/src/segcore/load_index_c.cpp @@ -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(info_proto.field().data_type()); + load_index_info->element_type = + static_cast(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(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(info_proto->field().data_type()); - load_index_info->element_type = static_cast( - 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(); diff --git a/internal/core/src/segcore/load_index_c.h b/internal/core/src/segcore/load_index_c.h index f69c3e94a0..e6e16a20d3 100644 --- a/internal/core/src/segcore/load_index_c.h +++ b/internal/core/src/segcore/load_index_c.h @@ -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); diff --git a/internal/proxy/task.go b/internal/proxy/task.go index dc19869a2c..8ce5d68f88 100644 --- a/internal/proxy/task.go +++ b/internal/proxy/task.go @@ -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 } diff --git a/internal/proxy/task_test.go b/internal/proxy/task_test.go index 0833f76ba5..72fad6f472 100644 --- a/internal/proxy/task_test.go +++ b/internal/proxy/task_test.go @@ -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() diff --git a/internal/querycoordv2/autoscale/resource.go b/internal/querycoordv2/autoscale/resource.go new file mode 100644 index 0000000000..880435363c --- /dev/null +++ b/internal/querycoordv2/autoscale/resource.go @@ -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 := ¶mtable.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 +} diff --git a/internal/querycoordv2/autoscale/resource_test.go b/internal/querycoordv2/autoscale/resource_test.go new file mode 100644 index 0000000000..7df63706fa --- /dev/null +++ b/internal/querycoordv2/autoscale/resource_test.go @@ -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"}, + }, + }, + }, + } +} diff --git a/internal/querycoordv2/autoscale_precheck.go b/internal/querycoordv2/autoscale_precheck.go new file mode 100644 index 0000000000..03ba299801 --- /dev/null +++ b/internal/querycoordv2/autoscale_precheck.go @@ -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) +} diff --git a/internal/querycoordv2/autoscale_precheck_test.go b/internal/querycoordv2/autoscale_precheck_test.go new file mode 100644 index 0000000000..e46e7057f4 --- /dev/null +++ b/internal/querycoordv2/autoscale_precheck_test.go @@ -0,0 +1,1271 @@ +// 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" + "testing" + "time" + + "github.com/cockroachdb/errors" + "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/samber/lo" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" + "github.com/stretchr/testify/require" + + "github.com/milvus-io/milvus-proto/go-api/v3/commonpb" + "github.com/milvus-io/milvus-proto/go-api/v3/milvuspb" + "github.com/milvus-io/milvus-proto/go-api/v3/schemapb" + metastoremocks "github.com/milvus-io/milvus/internal/metastore/mocks" + "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/storage" + "github.com/milvus-io/milvus/internal/util/sessionutil" + "github.com/milvus-io/milvus/pkg/v3/common" + "github.com/milvus-io/milvus/pkg/v3/metrics" + "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" +) + +func setAutoscaleCapacityRatiosForTest(t *testing.T, memoryHighWaterLevel, maxDiskUsagePercentage string) { + t.Helper() + params := paramtable.Get() + require.NoError(t, params.Save(params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key, memoryHighWaterLevel)) + require.NoError(t, params.Save(params.QueryNodeCfg.MaxDiskUsagePercentage.Key, maxDiskUsagePercentage)) + t.Cleanup(func() { + params.Reset(params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key) + params.Reset(params.QueryNodeCfg.MaxDiskUsagePercentage.Key) + }) +} + +func TestAutoscalePrecheckDecision(t *testing.T) { + setAutoscaleCapacityRatiosForTest(t, "1", "100") + + required := autoscaleResourceUsage{memoryBytes: 200, diskBytes: 160} + current := autoscaleResourceUsage{memoryBytes: 100, diskBytes: 100} + globalCapacity := autoscaleResourceUsage{memoryBytes: 900, diskBytes: 940} + check := func(required, available, globalCapacity, globalLimit autoscaleResourceUsage, autoscaleEnabled bool) error { + rgName := meta.DefaultResourceGroupName + return checkAutoscaleLoadResourceByRGWithCapacity( + map[string]autoscaleResourceUsage{rgName: required}, + map[string]autoscaleResourceUsage{rgName: available}, + map[string]autoscaleResourceUsage{rgName: available}, + globalCapacity, + globalLimit, + autoscaleEnabled, + ) + } + + assert.NoError(t, check(required, required, globalCapacity, autoscaleResourceUsage{}, false)) + assert.Error(t, check(required, current, globalCapacity, autoscaleResourceUsage{}, false)) + assert.Error(t, check(required, current, globalCapacity, autoscaleResourceUsage{}, true)) + assert.NoError(t, check(required, current, globalCapacity, autoscaleResourceUsage{memoryBytes: 1000, diskBytes: 1000}, true)) + assert.NoError(t, check( + autoscaleResourceUsage{memoryBytes: 200, diskBytes: 100}, + current, + globalCapacity, + autoscaleResourceUsage{memoryBytes: 1000}, + true, + )) + + err := check(required, current, globalCapacity, autoscaleResourceUsage{memoryBytes: 999, diskBytes: 1000}, true) + require.Error(t, err) + + var precheckErr *loadResourcePrecheckError + require.True(t, errors.As(err, &precheckErr)) + assert.Equal(t, int64(100), precheckErr.suggestedExpandPercent) + assert.ErrorIs(t, err, merr.ErrServiceResourceInsufficient) + assert.Equal(t, "Insufficient cluster resources. Please scale out the cluster or enable autoscaling. Suggested scale-out ratio: 100%", err.Error()) +} + +func TestAutoscalePrecheckDecisionByResourceGroup(t *testing.T) { + setAutoscaleCapacityRatiosForTest(t, "1", "100") + + requiredByRG := map[string]autoscaleResourceUsage{ + "rg-a": {memoryBytes: 200, diskBytes: 100}, + "rg-b": {memoryBytes: 100, diskBytes: 100}, + } + currentByRG := map[string]autoscaleResourceUsage{ + "rg-a": {memoryBytes: 100, diskBytes: 100}, + "rg-b": {memoryBytes: 1000, diskBytes: 1000}, + } + + err := checkAutoscaleLoadResourceByRGWithCapacity(requiredByRG, currentByRG, currentByRG, autoscaleResourceUsage{}, autoscaleResourceUsage{}, false) + require.Error(t, err) + + var precheckErr *loadResourcePrecheckError + require.True(t, errors.As(err, &precheckErr)) + assert.Equal(t, int64(100), precheckErr.suggestedExpandPercent) + + assert.NoError(t, checkAutoscaleLoadResourceByRGWithCapacity( + requiredByRG, + currentByRG, + currentByRG, + autoscaleResourceUsage{memoryBytes: 900}, + autoscaleResourceUsage{memoryBytes: 1000}, + true, + )) + assert.Error(t, checkAutoscaleLoadResourceByRGWithCapacity( + requiredByRG, + currentByRG, + currentByRG, + autoscaleResourceUsage{memoryBytes: 901}, + autoscaleResourceUsage{memoryBytes: 1000}, + true, + )) +} + +func TestAutoscalePrecheckGlobalLimitUsesTotalCapacity(t *testing.T) { + setAutoscaleCapacityRatiosForTest(t, "0.9", "90") + + globalCapacity := autoscaleResourceUsage{memoryBytes: 80, diskBytes: 80} + globalLimit := autoscaleResourceUsage{memoryBytes: 100, diskBytes: 100} + + assert.True(t, canExpandWithinGlobalLimit( + autoscaleResourceUsage{memoryBytes: 18, diskBytes: 18}, + globalCapacity, + globalLimit, + )) + assert.False(t, canExpandWithinGlobalLimit( + autoscaleResourceUsage{memoryBytes: 19, diskBytes: 18}, + globalCapacity, + globalLimit, + )) + assert.False(t, canExpandWithinGlobalLimit( + autoscaleResourceUsage{memoryBytes: 18, diskBytes: 19}, + globalCapacity, + globalLimit, + )) + assert.False(t, canExpandWithinGlobalLimit( + autoscaleResourceUsage{memoryBytes: 1, diskBytes: 1}, + globalLimit, + globalLimit, + )) +} + +func TestAutoscalePrecheckSuggestionUsesCapacityBase(t *testing.T) { + requiredByRG := map[string]autoscaleResourceUsage{ + "rg-a": {memoryBytes: 200, diskBytes: 160}, + } + availableByRG := map[string]autoscaleResourceUsage{ + "rg-a": {memoryBytes: 100, diskBytes: 100}, + } + capacityByRG := map[string]autoscaleResourceUsage{ + "rg-a": {memoryBytes: 1000, diskBytes: 1000}, + } + + err := checkAutoscaleLoadResourceByRGWithCapacity( + requiredByRG, + availableByRG, + capacityByRG, + autoscaleResourceUsage{}, + autoscaleResourceUsage{}, + false, + ) + require.Error(t, err) + + var precheckErr *loadResourcePrecheckError + require.True(t, errors.As(err, &precheckErr)) + assert.Equal(t, int64(10), precheckErr.suggestedExpandPercent) +} + +func TestAutoscalePrecheckIncludesNegativeAvailableResourceInShortage(t *testing.T) { + err := checkAutoscaleLoadResourceByRGWithCapacity( + map[string]autoscaleResourceUsage{"rg": {memoryBytes: 10}}, + map[string]autoscaleResourceUsage{"rg": {memoryBytes: -20}}, + map[string]autoscaleResourceUsage{"rg": {memoryBytes: 100}}, + autoscaleResourceUsage{}, + autoscaleResourceUsage{}, + false, + ) + + require.Error(t, err) + var precheckErr *loadResourcePrecheckError + require.ErrorAs(t, err, &precheckErr) + assert.Equal(t, int64(30), precheckErr.suggestedExpandPercent) +} + +func TestAutoscalePrecheckGlobalLimitUsesGB(t *testing.T) { + limit := autoscaleGlobalLimitFromConfig(2, 3) + + assert.Equal(t, int64(2*1024*1024*1024), limit.memoryBytes) + assert.Equal(t, int64(3*1024*1024*1024), limit.diskBytes) +} + +func TestBuildRequiredLoadResourceByRGForLoadConfig(t *testing.T) { + requiredByRG := buildRequiredLoadResourceByRGForLoadConfig( + autoscaleResourceUsage{memoryBytes: 300, diskBytes: 200}, + autoscaleResourceUsage{memoryBytes: 100, diskBytes: 80}, + map[string]int{ + "rg-a": 1, + "rg-b": 1, + }, + map[string]int{ + "rg-a": 2, + "rg-b": 0, + "rg-c": 1, + }, + ) + + assert.Equal(t, map[string]autoscaleResourceUsage{ + "rg-a": {memoryBytes: 500, diskBytes: 320}, + "rg-c": {memoryBytes: 300, diskBytes: 200}, + }, requiredByRG) +} + +func TestAutoscaleLoadResourceFiltersLoadFieldsAndKeepsAllIndexesPerField(t *testing.T) { + segments := []*datapb.SegmentInfo{ + { + ID: 1, + PartitionID: 10, + Binlogs: []*datapb.FieldBinlog{ + {FieldID: 100}, + {FieldID: 101}, + }, + Statslogs: []*datapb.FieldBinlog{ + {FieldID: 100}, + {FieldID: 101}, + }, + Deltalogs: []*datapb.FieldBinlog{ + {FieldID: 999}, + }, + Bm25Statslogs: []*datapb.FieldBinlog{ + {FieldID: 100}, + {FieldID: 101}, + }, + JsonKeyStats: map[int64]*datapb.JsonKeyStats{ + 100: {}, + 101: {}, + }, + TextStatsLogs: map[int64]*datapb.TextIndexStats{ + 100: {}, + 101: {}, + }, + }, + } + + filteredSegments := autoscale.FilterSegmentsByLoadFields(segments, []int64{100}) + require.Len(t, filteredSegments, 1) + require.Len(t, filteredSegments[0].GetBinlogs(), 1) + assert.Equal(t, int64(100), filteredSegments[0].GetBinlogs()[0].GetFieldID()) + require.Len(t, filteredSegments[0].GetStatslogs(), 1) + assert.Equal(t, int64(100), filteredSegments[0].GetStatslogs()[0].GetFieldID()) + require.Len(t, filteredSegments[0].GetBm25Statslogs(), 1) + assert.Equal(t, int64(100), filteredSegments[0].GetBm25Statslogs()[0].GetFieldID()) + assert.Contains(t, filteredSegments[0].GetJsonKeyStats(), int64(100)) + assert.NotContains(t, filteredSegments[0].GetJsonKeyStats(), int64(101)) + assert.Contains(t, filteredSegments[0].GetTextStatsLogs(), int64(100)) + assert.NotContains(t, filteredSegments[0].GetTextStatsLogs(), int64(101)) + require.Len(t, filteredSegments[0].GetDeltalogs(), 1) + require.Len(t, segments[0].GetBinlogs(), 2) + + indexes := map[int64][]*querypb.FieldIndexInfo{ + 1: { + {FieldID: 100, IndexID: 10, IndexParams: []*commonpb.KeyValuePair{{Key: common.JSONPathKey, Value: "/a"}}}, + {FieldID: 100, IndexID: 11, IndexParams: []*commonpb.KeyValuePair{{Key: common.JSONPathKey, Value: "/b"}}}, + {FieldID: 101, IndexID: 20}, + }, + } + + filteredIndexes := autoscale.FilterIndexesByLoadFields(indexes, []int64{100}) + require.Len(t, filteredIndexes[1], 2) + assert.ElementsMatch(t, []int64{10, 11}, lo.Map(filteredIndexes[1], func(index *querypb.FieldIndexInfo, _ int) int64 { + return index.GetIndexID() + })) +} + +func TestGetIndexInfoBySegmentsBatchesRequests(t *testing.T) { + ctx := context.Background() + collectionID := int64(100) + segmentIDs := make([]int64, maxSegmentNumPerGetIndexInfoRequest+1) + for i := range segmentIDs { + segmentIDs[i] = int64(i + 1) + } + + firstBatchArgs := make([]interface{}, 0, maxSegmentNumPerGetIndexInfoRequest) + for _, segmentID := range segmentIDs[:maxSegmentNumPerGetIndexInfoRequest] { + firstBatchArgs = append(firstBatchArgs, segmentID) + } + broker := meta.NewMockBroker(t) + broker.EXPECT(). + GetIndexInfo(mock.Anything, collectionID, firstBatchArgs...). + Return(map[int64][]*querypb.FieldIndexInfo{ + segmentIDs[0]: {{FieldID: 10, IndexID: 100}}, + }, nil). + Once() + broker.EXPECT(). + GetIndexInfo(mock.Anything, collectionID, segmentIDs[maxSegmentNumPerGetIndexInfoRequest]). + Return(map[int64][]*querypb.FieldIndexInfo{ + segmentIDs[maxSegmentNumPerGetIndexInfoRequest]: {{FieldID: 11, IndexID: 101}}, + }, nil). + Once() + + indexes, err := getIndexInfoBySegments(ctx, broker, collectionID, segmentIDs) + + require.NoError(t, err) + require.Len(t, indexes, 2) + assert.Equal(t, int64(100), indexes[segmentIDs[0]][0].GetIndexID()) + assert.Equal(t, int64(101), indexes[segmentIDs[maxSegmentNumPerGetIndexInfoRequest]][0].GetIndexID()) +} + +func TestMemoryResourceFromMetricsUsesQueryNodeMetrics(t *testing.T) { + params := paramtable.Get() + params.Save(params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key, "0.4") + defer params.Reset(params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key) + params.Save(params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key, "0.5") + defer params.Reset(params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key) + + for _, test := range []struct { + usedMemory uint64 + expectedAvailable int64 + }{ + {usedMemory: 400, expectedAvailable: 100}, + {usedMemory: 499, expectedAvailable: 1}, + {usedMemory: 500, expectedAvailable: 0}, + {usedMemory: 600, expectedAvailable: -100}, + } { + availableMemoryBytes, capacityMemoryBytes := memoryResourceFromMetrics(1000, test.usedMemory) + assert.Equal(t, test.expectedAvailable, availableMemoryBytes) + assert.Equal(t, int64(1000), capacityMemoryBytes) + } +} + +const ( + autoscalePrecheckTestCollectionID = int64(100) + autoscalePrecheckTestPartitionID = int64(10) + autoscalePrecheckTestSegmentID = int64(1) +) + +func newAutoscalePrecheckTestRequest() *job.AlterLoadConfigRequest { + return &job.AlterLoadConfigRequest{ + CollectionInfo: &milvuspb.DescribeCollectionResponse{ + CollectionID: autoscalePrecheckTestCollectionID, + Schema: &schemapb.CollectionSchema{ + Fields: []*schemapb.FieldSchema{{FieldID: 100, Name: "field_100", DataType: schemapb.DataType_Int64}}, + }, + }, + Expected: job.ExpectedLoadConfig{ + ExpectedPartitionIDs: []int64{autoscalePrecheckTestPartitionID}, + ExpectedReplicaNumber: map[string]int{meta.DefaultResourceGroupName: 1}, + }, + } +} + +func TestRequiredLoadResourceAppliesCollectionMmap(t *testing.T) { + params := paramtable.Get() + require.NoError(t, params.Save(params.QueryNodeCfg.TieredEvictionEnabled.Key, "false")) + t.Cleanup(func() { + params.Reset(params.QueryNodeCfg.TieredEvictionEnabled.Key) + params.Reset(params.QueryNodeCfg.MmapScalarField.Key) + }) + for _, test := range []struct { + name, global, collection, field string + wantDisk bool + }{ + {"collection enables mmap", "false", "true", "", true}, + {"collection disables mmap", "true", "false", "", false}, + {"field disables mmap", "false", "true", "false", false}, + {"field enables mmap", "true", "false", "true", true}, + {"global default", "true", "", "", true}, + } { + t.Run(test.name, func(t *testing.T) { + require.NoError(t, params.Save(params.QueryNodeCfg.MmapScalarField.Key, test.global)) + req := newAutoscalePrecheckTestRequest() + if test.collection != "" { + req.CollectionInfo.Properties = []*commonpb.KeyValuePair{{Key: common.MmapEnabledKey, Value: test.collection}} + } + field := req.CollectionInfo.Schema.Fields[0] + if test.field != "" { + field.TypeParams = []*commonpb.KeyValuePair{{Key: common.MmapEnabledKey, Value: test.field}} + } + server := &Server{broker: newAutoscalePrecheckTestBroker(t, nil, nil, nil)} + + usage, err := server.requiredLoadResourceByRG(context.Background(), req) + + require.NoError(t, err) + want := autoscaleResourceUsage{memoryBytes: 1024} + if test.wantDisk { + want = autoscaleResourceUsage{diskBytes: 1024} + } + assert.Equal(t, want, usage[meta.DefaultResourceGroupName]) + if test.field == "" { + assert.Empty(t, field.TypeParams, "estimation must not mutate the source schema") + } else { + assert.Equal(t, []*commonpb.KeyValuePair{{Key: common.MmapEnabledKey, Value: test.field}}, field.TypeParams) + } + }) + } +} + +func enableAutoscalePrecheckForTest(t *testing.T) { + t.Helper() + params := paramtable.Get() + params.Save(params.QueryCoordCfg.AutoscalePrecheckEnabled.Key, "true") + params.Save(params.QueryNodeCfg.TieredEvictionEnabled.Key, "false") + t.Cleanup(func() { + params.Reset(params.QueryCoordCfg.AutoscalePrecheckEnabled.Key) + params.Reset(params.QueryNodeCfg.TieredEvictionEnabled.Key) + }) +} + +func newAutoscalePrecheckTestBroker(t *testing.T, recoveryErr, segmentErr, indexErr error) *meta.MockBroker { + t.Helper() + broker := meta.NewMockBroker(t) + if recoveryErr != nil { + broker.EXPECT(). + GetRecoveryInfoV2(mock.Anything, autoscalePrecheckTestCollectionID, autoscalePrecheckTestPartitionID). + Return(nil, nil, recoveryErr). + Once() + return broker + } + + broker.EXPECT(). + GetRecoveryInfoV2(mock.Anything, autoscalePrecheckTestCollectionID, autoscalePrecheckTestPartitionID). + Return(nil, []*datapb.SegmentInfo{ + { + ID: autoscalePrecheckTestSegmentID, + PartitionID: autoscalePrecheckTestPartitionID, + }, + }, nil). + Once() + broker.EXPECT(). + GetSegmentInfo(mock.Anything, autoscalePrecheckTestSegmentID). + Return([]*datapb.SegmentInfo{ + { + ID: autoscalePrecheckTestSegmentID, + PartitionID: autoscalePrecheckTestPartitionID, + NumOfRows: 1, + StorageVersion: storage.StorageV3, + ManifestPath: "files/manifest", + Stats: &datapb.Statistics{ + LoadResource: &datapb.LoadResourceStatistics{ + ColumnGroups: []*datapb.ColumnGroupStatistics{{ + GroupId: 100, + FieldIds: []int64{100}, + MemorySize: 1024, + }}, + }, + }, + }, + }, segmentErr). + Once() + if segmentErr != nil { + return broker + } + broker.EXPECT(). + GetIndexInfo(mock.Anything, autoscalePrecheckTestCollectionID, autoscalePrecheckTestSegmentID). + Return(nil, indexErr). + Once() + return broker +} + +func newAutoscalePrecheckTestServer(t *testing.T, cluster session.Cluster) *Server { + t.Helper() + nodeManager := session.NewNodeManager() + nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1})) + catalog := metastoremocks.NewQueryCoordCatalog(t) + catalog.EXPECT(). + SaveResourceGroup(mock.Anything, mock.Anything). + Return(nil). + Once() + metaManager := meta.NewMeta(nil, catalog, nodeManager) + metaManager.HandleNodeUp(context.Background(), 1) + return &Server{ + meta: metaManager, + nodeMgr: nodeManager, + cluster: cluster, + } +} + +func newAutoscalePrecheckMetrics(t *testing.T, memory, memoryUsage uint64, disk float64) string { + t.Helper() + localDiskCapacity := int64(disk * 1e9) + metrics, err := metricsinfo.MarshalComponentInfos(metricsinfo.QueryNodeInfos{ + BaseComponentInfos: metricsinfo.BaseComponentInfos{ + HardwareInfos: metricsinfo.HardwareMetrics{ + Memory: memory, + MemoryUsage: memoryUsage, + Disk: disk, + }, + }, + LocalStorage: &metricsinfo.QueryNodeLocalStorageMetrics{ + CapacityBytes: localDiskCapacity, + UsedBytes: 0, + }, + }) + require.NoError(t, err) + return metrics +} + +func TestCurrentQueryNodeResourceStatusByRGUsesReportedAvailability(t *testing.T) { + params := paramtable.Get() + params.Save(params.QueryCoordCfg.AutoscaleEnabled.Key, "false") + defer params.Reset(params.QueryCoordCfg.AutoscaleEnabled.Key) + + metricsPayload := newAutoscalePrecheckMetrics(t, 1024, 512, 1) + cluster := session.NewMockCluster(t) + cluster.EXPECT().GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil). + Once() + server := newAutoscalePrecheckTestServer(t, cluster) + + rawMemory, capacityMemory := memoryResourceFromMetrics(1024, 512) + usableCapacity := usableResourceCapacity(autoscaleResourceUsage{memoryBytes: capacityMemory, diskBytes: 1e9}) + availableByRG, capacityByRG, _, err := server.currentQueryNodeResourceStatus( + context.Background(), + map[string]autoscaleResourceUsage{meta.DefaultResourceGroupName: {memoryBytes: 1}}, + false, + ) + + require.NoError(t, err) + assert.Equal(t, rawMemory, availableByRG[meta.DefaultResourceGroupName].memoryBytes) + assert.Equal(t, usableCapacity.diskBytes, availableByRG[meta.DefaultResourceGroupName].diskBytes) + assert.Equal(t, usableCapacity.memoryBytes, capacityByRG[meta.DefaultResourceGroupName].memoryBytes) + assert.Equal(t, usableCapacity.diskBytes, capacityByRG[meta.DefaultResourceGroupName].diskBytes) +} + +func TestCurrentQueryNodeResourceStatusUsesQueryNodeDiskQuota(t *testing.T) { + params := paramtable.Get() + params.Save(params.QueryNodeCfg.MaxDiskUsagePercentage.Key, "95") + defer params.Reset(params.QueryNodeCfg.MaxDiskUsagePercentage.Key) + + const gib = int64(1024 * 1024 * 1024) + filesystemBytes := int64(3584) * gib + workerUsableBytes := int64(1520) * gib + workerUsedBytes := 17 * gib / 10 + workerAvailableBytes := workerUsableBytes - workerUsedBytes + physicalUsedBytes := 1312 * gib / 10 + metricsPayload, err := metricsinfo.MarshalComponentInfos(metricsinfo.QueryNodeInfos{ + BaseComponentInfos: metricsinfo.BaseComponentInfos{ + HardwareInfos: metricsinfo.HardwareMetrics{ + Memory: 1024, + MemoryUsage: 512, + Disk: float64(filesystemBytes) / 1e9, + DiskUsage: float64(physicalUsedBytes) / 1e9, + }, + }, + LocalStorage: &metricsinfo.QueryNodeLocalStorageMetrics{ + CapacityBytes: 1600 * gib, + UsedBytes: workerUsedBytes, + }, + }) + require.NoError(t, err) + cluster := session.NewMockCluster(t) + cluster.EXPECT().GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil).Once() + server := newAutoscalePrecheckTestServer(t, cluster) + + availableByRG, capacityByRG, globalCapacity, err := server.currentQueryNodeResourceStatus( + context.Background(), + map[string]autoscaleResourceUsage{meta.DefaultResourceGroupName: {diskBytes: 1}}, + true, + ) + + require.NoError(t, err) + assert.Equal(t, workerAvailableBytes, availableByRG[meta.DefaultResourceGroupName].diskBytes) + assert.Equal(t, workerUsableBytes, capacityByRG[meta.DefaultResourceGroupName].diskBytes) + assert.Equal(t, int64(1600)*gib, globalCapacity.diskBytes) + assert.True(t, canExpandWithinGlobalLimit( + autoscaleResourceUsage{diskBytes: 1000 * gib}, + globalCapacity, + autoscaleResourceUsage{diskBytes: 3000 * gib}, + )) + assert.False(t, canExpandWithinGlobalLimit( + autoscaleResourceUsage{diskBytes: 1000 * gib}, + globalCapacity, + autoscaleResourceUsage{diskBytes: 2000 * gib}, + )) +} + +func TestCurrentQueryNodeResourceStatusUsesLocalBudgetWithoutPhysicalMinimum(t *testing.T) { + params := paramtable.Get() + params.Save(params.QueryNodeCfg.MaxDiskUsagePercentage.Key, "95") + defer params.Reset(params.QueryNodeCfg.MaxDiskUsagePercentage.Key) + + const gib = int64(1024 * 1024 * 1024) + const filesystemBytes = int64(1000) * gib + physicalUsedBytes := int64(900) * gib + metricsPayload, err := metricsinfo.MarshalComponentInfos(metricsinfo.QueryNodeInfos{ + BaseComponentInfos: metricsinfo.BaseComponentInfos{ + HardwareInfos: metricsinfo.HardwareMetrics{ + Memory: 1024, + MemoryUsage: 512, + Disk: float64(filesystemBytes) / 1e9, + DiskUsage: float64(physicalUsedBytes) / 1e9, + }, + }, + LocalStorage: &metricsinfo.QueryNodeLocalStorageMetrics{ + CapacityBytes: 1600 * gib, + UsedBytes: 2 * gib, + }, + }) + require.NoError(t, err) + cluster := session.NewMockCluster(t) + cluster.EXPECT().GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil).Once() + server := newAutoscalePrecheckTestServer(t, cluster) + + availableByRG, capacityByRG, _, err := server.currentQueryNodeResourceStatus( + context.Background(), + map[string]autoscaleResourceUsage{meta.DefaultResourceGroupName: {diskBytes: 1}}, + false, + ) + + require.NoError(t, err) + assert.Equal(t, int64(1518)*gib, availableByRG[meta.DefaultResourceGroupName].diskBytes) + assert.Equal(t, int64(1520)*gib, capacityByRG[meta.DefaultResourceGroupName].diskBytes) +} + +func TestCurrentQueryNodeResourceStatusDoesNotRequireFilesystemDiskMetrics(t *testing.T) { + const gib = int64(1024 * 1024 * 1024) + metricsPayload, err := metricsinfo.MarshalComponentInfos(metricsinfo.QueryNodeInfos{ + BaseComponentInfos: metricsinfo.BaseComponentInfos{ + HardwareInfos: metricsinfo.HardwareMetrics{ + Memory: 1024, + MemoryUsage: 512, + }, + }, + LocalStorage: &metricsinfo.QueryNodeLocalStorageMetrics{ + CapacityBytes: 1600 * gib, + UsedBytes: 2 * gib, + }, + }) + require.NoError(t, err) + cluster := session.NewMockCluster(t) + cluster.EXPECT().GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil).Once() + server := newAutoscalePrecheckTestServer(t, cluster) + + availableByRG, capacityByRG, globalCapacity, err := server.currentQueryNodeResourceStatus( + context.Background(), + map[string]autoscaleResourceUsage{meta.DefaultResourceGroupName: {diskBytes: 1}}, + true, + ) + + require.NoError(t, err) + usableDiskBytes := usableResourceCapacity(autoscaleResourceUsage{diskBytes: 1600 * gib}).diskBytes + assert.Equal(t, usableDiskBytes-2*gib, availableByRG[meta.DefaultResourceGroupName].diskBytes) + assert.Equal(t, usableDiskBytes, capacityByRG[meta.DefaultResourceGroupName].diskBytes) + assert.Equal(t, int64(1600)*gib, globalCapacity.diskBytes) +} + +func TestQueryNodeResourceSnapshotSharedFilesystemUsesLocalStorageQuotas(t *testing.T) { + params := paramtable.Get() + require.NoError(t, params.Save(params.QueryNodeCfg.MaxDiskUsagePercentage.Key, "95")) + t.Cleanup(func() { params.Reset(params.QueryNodeCfg.MaxDiskUsagePercentage.Key) }) + + const gib = int64(1024 * 1024 * 1024) + metricsPayload, err := metricsinfo.MarshalComponentInfos(metricsinfo.QueryNodeInfos{ + BaseComponentInfos: metricsinfo.BaseComponentInfos{ + HardwareInfos: metricsinfo.HardwareMetrics{ + Memory: 1024, + MemoryUsage: 512, + Disk: 1000, + DiskUsage: 500, + }, + }, + LocalStorage: &metricsinfo.QueryNodeLocalStorageMetrics{ + CapacityBytes: 300 * gib, + UsedBytes: 10 * gib, + }, + }) + require.NoError(t, err) + + cluster := session.NewMockCluster(t) + nodes := make([]*session.NodeInfo, 0, 3) + for nodeID := int64(1); nodeID <= 3; nodeID++ { + nodes = append(nodes, session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: nodeID})) + cluster.EXPECT().GetMetrics(mock.Anything, nodeID, mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil).Once() + } + server := &Server{cluster: cluster} + + snapshot, err := server.queryNodeResourceSnapshot(context.Background(), nodes) + require.NoError(t, err) + available, capacity, ok := queryNodeResourceStatusFromSnapshot(nodes, snapshot) + require.True(t, ok) + assert.Equal(t, int64(3*275)*gib, available.diskBytes) + assert.Equal(t, int64(3*300)*gib, capacity.diskBytes) +} + +func TestCurrentQueryNodeResourceStatusByRGUsesActiveResourceGroupNodes(t *testing.T) { + params := paramtable.Get() + params.Save(params.QueryCoordCfg.AutoscaleEnabled.Key, "false") + defer params.Reset(params.QueryCoordCfg.AutoscaleEnabled.Key) + + const ( + node1Host = "task-2-active-rg-node-1" + node2Host = "task-2-stopping-rg-node-2" + node3Host = "task-2-online-outsider-node-3" + ) + nodeManager := session.NewNodeManager() + nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1, Hostname: node1Host})) + nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 2, Hostname: node2Host})) + nodeManager.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 3, Hostname: node3Host})) + catalog := metastoremocks.NewQueryCoordCatalog(t) + catalog.EXPECT(). + SaveResourceGroup(mock.Anything, mock.Anything). + Return(nil). + Twice() + metaManager := meta.NewMeta(nil, catalog, nodeManager) + metaManager.HandleNodeUp(context.Background(), 1) + metaManager.HandleNodeUp(context.Background(), 2) + nodeManager.Stopping(2) + + metricsPayload := newAutoscalePrecheckMetrics(t, 1024, 512, 1) + cluster := session.NewMockCluster(t) + cluster.EXPECT(). + GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil). + Once() + server := &Server{ + meta: metaManager, + nodeMgr: nodeManager, + cluster: cluster, + } + + rawMemory, capacityMemory := memoryResourceFromMetrics(1024, 512) + usableCapacity := usableResourceCapacity(autoscaleResourceUsage{memoryBytes: capacityMemory, diskBytes: 1e9}) + availableByRG, capacityByRG, _, err := server.currentQueryNodeResourceStatus( + context.Background(), + map[string]autoscaleResourceUsage{meta.DefaultResourceGroupName: {memoryBytes: 1}}, + false, + ) + + require.NoError(t, err) + assert.Equal(t, rawMemory, availableByRG[meta.DefaultResourceGroupName].memoryBytes) + assert.Equal(t, usableCapacity.diskBytes, availableByRG[meta.DefaultResourceGroupName].diskBytes) + assert.Equal(t, usableCapacity.memoryBytes, capacityByRG[meta.DefaultResourceGroupName].memoryBytes) + assert.Equal(t, usableCapacity.diskBytes, capacityByRG[meta.DefaultResourceGroupName].diskBytes) +} + +func TestCheckLoadResourceSkipsEstimateDependencyErrors(t *testing.T) { + enableAutoscalePrecheckForTest(t) + + tests := []struct { + name string + recoveryErr error + segmentErr error + indexErr error + }{ + { + name: "recovery info error", + recoveryErr: errors.New("recovery info unavailable"), + }, + { + name: "segment info error", + segmentErr: errors.New("segment info unavailable"), + }, + { + name: "index info error", + indexErr: errors.New("index info unavailable"), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + server := &Server{broker: newAutoscalePrecheckTestBroker(t, test.recoveryErr, test.segmentErr, test.indexErr)} + + requiredByRG, err := server.checkLoadResource(context.Background(), newAutoscalePrecheckTestRequest()) + + require.NoError(t, err) + assert.Empty(t, requiredByRG) + }) + } +} + +func TestCheckLoadResourceSkipsEstimateError(t *testing.T) { + enableAutoscalePrecheckForTest(t) + + req := newAutoscalePrecheckTestRequest() + req.CollectionInfo.Schema = nil + server := &Server{broker: newAutoscalePrecheckTestBroker(t, nil, nil, nil)} + memoryCounter := metrics.QueryCoordLoadDemandMemoryBytes.WithLabelValues(meta.DefaultResourceGroupName) + diskCounter := metrics.QueryCoordLoadDemandDiskBytes.WithLabelValues(meta.DefaultResourceGroupName) + memoryBefore := testutil.ToFloat64(memoryCounter) + diskBefore := testutil.ToFloat64(diskCounter) + + requiredByRG, err := server.checkLoadResource(context.Background(), req) + require.NoError(t, err) + assert.Empty(t, requiredByRG) + assert.Equal(t, memoryBefore, testutil.ToFloat64(memoryCounter)) + assert.Equal(t, diskBefore, testutil.ToFloat64(diskCounter)) +} + +func TestCheckLoadResourceDisabledSkipsPrecheck(t *testing.T) { + params := paramtable.Get() + params.Save(params.QueryCoordCfg.AutoscalePrecheckEnabled.Key, "false") + t.Cleanup(func() { + params.Reset(params.QueryCoordCfg.AutoscalePrecheckEnabled.Key) + }) + + server := &Server{} + requiredByRG, err := server.checkLoadResource(context.Background(), newAutoscalePrecheckTestRequest()) + require.NoError(t, err) + assert.Empty(t, requiredByRG) +} + +func TestRecordLoadResourceDemand(t *testing.T) { + const rgName = "load-demand-metric-test" + metrics.QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(rgName) + metrics.QueryCoordLoadDemandDiskBytes.DeleteLabelValues(rgName) + memoryCounter := metrics.QueryCoordLoadDemandMemoryBytes.WithLabelValues(rgName) + diskCounter := metrics.QueryCoordLoadDemandDiskBytes.WithLabelValues(rgName) + t.Cleanup(func() { + metrics.QueryCoordLoadDemandMemoryBytes.DeleteLabelValues(rgName) + metrics.QueryCoordLoadDemandDiskBytes.DeleteLabelValues(rgName) + }) + + recordLoadResourceDemand(map[string]autoscaleResourceUsage{ + rgName: {memoryBytes: 100, diskBytes: 80}, + }) + recordLoadResourceDemand(map[string]autoscaleResourceUsage{ + rgName: {memoryBytes: 20}, + }) + + assert.Equal(t, float64(120), testutil.ToFloat64(memoryCounter)) + assert.Equal(t, float64(80), testutil.ToFloat64(diskCounter)) +} + +func TestCheckLoadResourceSkipsResourceGroupMetricsErrors(t *testing.T) { + enableAutoscalePrecheckForTest(t) + + params := paramtable.Get() + params.Save(params.QueryCoordCfg.AutoscaleEnabled.Key, "false") + defer params.Reset(params.QueryCoordCfg.AutoscaleEnabled.Key) + + validMetrics := newAutoscalePrecheckMetrics(t, 1024, 512, 1) + tests := []struct { + name string + resp *milvuspb.GetMetricsResponse + err error + }{ + { + name: "rpc error", + err: errors.New("metrics unavailable"), + }, + { + name: "nil response", + }, + { + name: "status error", + resp: &milvuspb.GetMetricsResponse{ + Status: merr.Status(merr.ErrServiceNotReady), + Response: validMetrics, + }, + }, + { + name: "malformed payload", + resp: &milvuspb.GetMetricsResponse{ + Status: merr.Success(), + Response: "{", + }, + }, + { + name: "missing memory metrics", + resp: &milvuspb.GetMetricsResponse{ + Status: merr.Success(), + Response: `{ + "hardware_infos": { + "disk": 1 + } + }`, + }, + }, + { + name: "missing local storage metrics", + resp: &milvuspb.GetMetricsResponse{ + Status: merr.Success(), + Response: `{ + "hardware_infos": { + "memory": 1024, + "memory_usage": 512, + "disk": 1 + } + }`, + }, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + cluster := session.NewMockCluster(t) + cluster.EXPECT(). + GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(test.resp, test.err). + Once() + server := newAutoscalePrecheckTestServer(t, cluster) + server.broker = newAutoscalePrecheckTestBroker(t, nil, nil, nil) + + requiredByRG, err := server.checkLoadResource(context.Background(), newAutoscalePrecheckTestRequest()) + + require.NoError(t, err) + assert.Equal(t, map[string]autoscaleResourceUsage{ + meta.DefaultResourceGroupName: {memoryBytes: 1024}, + }, requiredByRG) + }) + } +} + +func TestCheckLoadResourceTieredPreservesDemandWithoutMetrics(t *testing.T) { + enableAutoscalePrecheckForTest(t) + params := paramtable.Get() + require.NoError(t, params.Save(params.QueryNodeCfg.TieredEvictionEnabled.Key, "true")) + require.NoError(t, params.Save(params.QueryNodeCfg.TieredEvictableMemoryCacheRatio.Key, "1")) + require.NoError(t, params.Save(params.QueryNodeCfg.MmapScalarField.Key, "false")) + t.Cleanup(func() { + params.Reset(params.QueryNodeCfg.TieredEvictableMemoryCacheRatio.Key) + params.Reset(params.QueryNodeCfg.MmapScalarField.Key) + }) + cluster := session.NewMockCluster(t) + server := newAutoscalePrecheckTestServer(t, cluster) + server.broker = newAutoscalePrecheckTestBroker(t, nil, nil, nil) + + requiredByRG, err := server.checkLoadResource(context.Background(), newAutoscalePrecheckTestRequest()) + + require.NoError(t, err) + assert.Equal(t, map[string]autoscaleResourceUsage{ + meta.DefaultResourceGroupName: {memoryBytes: 1024}, + }, requiredByRG) + cluster.AssertNotCalled(t, "GetMetrics", mock.Anything, mock.Anything, mock.Anything) +} + +func TestCheckLoadResourceSkipsMetricsTimeout(t *testing.T) { + enableAutoscalePrecheckForTest(t) + cluster := session.NewMockCluster(t) + cluster.EXPECT().GetMetrics(mock.Anything, int64(1), mock.Anything). + RunAndReturn(func(ctx context.Context, _ int64, _ *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) { + deadline, ok := ctx.Deadline() + assert.True(t, ok) + assert.LessOrEqual(t, time.Until(deadline), 10*time.Second) + <-ctx.Done() + return nil, ctx.Err() + }).Once() + server := newAutoscalePrecheckTestServer(t, cluster) + server.broker = newAutoscalePrecheckTestBroker(t, nil, nil, nil) + ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second) + defer cancel() + + requiredByRG, err := server.checkLoadResource(ctx, newAutoscalePrecheckTestRequest()) + + require.NoError(t, err) + assert.Equal(t, map[string]autoscaleResourceUsage{ + meta.DefaultResourceGroupName: {memoryBytes: 1024}, + }, requiredByRG) + assert.NoError(t, ctx.Err(), "metrics collection should time out before the load request") +} + +func TestCheckLoadResourceSkipsUnavailableResourceGroupNodes(t *testing.T) { + enableAutoscalePrecheckForTest(t) + params := paramtable.Get() + require.NoError(t, params.Save(params.QueryCoordCfg.AutoscaleEnabled.Key, "false")) + t.Cleanup(func() { params.Reset(params.QueryCoordCfg.AutoscaleEnabled.Key) }) + + for _, stopping := range []bool{false, true} { + name := "offline" + if stopping { + name = "stopping" + } + t.Run(name, func(t *testing.T) { + cluster := session.NewMockCluster(t) + server := newAutoscalePrecheckTestServer(t, cluster) + server.broker = newAutoscalePrecheckTestBroker(t, nil, nil, nil) + if stopping { + server.nodeMgr.Stopping(1) + } else { + server.nodeMgr.Remove(1) + } + + requiredByRG, err := server.checkLoadResource(context.Background(), newAutoscalePrecheckTestRequest()) + + require.NoError(t, err) + assert.Equal(t, map[string]autoscaleResourceUsage{ + meta.DefaultResourceGroupName: {memoryBytes: 1024}, + }, requiredByRG) + cluster.AssertNotCalled(t, "GetMetrics", mock.Anything, mock.Anything, mock.Anything) + }) + } +} + +func TestCheckLoadResourceSkipsGlobalMetricsError(t *testing.T) { + enableAutoscalePrecheckForTest(t) + + params := paramtable.Get() + params.Save(params.QueryCoordCfg.AutoscaleEnabled.Key, "true") + defer params.Reset(params.QueryCoordCfg.AutoscaleEnabled.Key) + params.Save(params.QueryCoordCfg.AutoscaleMaxMemoryLimit.Key, "1") + defer params.Reset(params.QueryCoordCfg.AutoscaleMaxMemoryLimit.Key) + params.Save(params.QueryCoordCfg.AutoscaleMaxDiskLimit.Key, "1") + defer params.Reset(params.QueryCoordCfg.AutoscaleMaxDiskLimit.Key) + + validMetrics := newAutoscalePrecheckMetrics(t, 1024, 1024, 1) + cluster := session.NewMockCluster(t) + cluster.EXPECT(). + GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: validMetrics}, nil). + Once() + cluster.EXPECT(). + GetMetrics(mock.Anything, int64(2), mock.Anything). + Return(nil, errors.New("global metrics unavailable")). + Once() + server := newAutoscalePrecheckTestServer(t, cluster) + server.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 2})) + server.broker = newAutoscalePrecheckTestBroker(t, nil, nil, nil) + + requiredByRG, err := server.checkLoadResource(context.Background(), newAutoscalePrecheckTestRequest()) + + require.NoError(t, err) + assert.Equal(t, map[string]autoscaleResourceUsage{ + meta.DefaultResourceGroupName: {memoryBytes: 1024}, + }, requiredByRG) +} + +func TestCheckLoadResourceReturnsAcceptedDemandWithoutRecording(t *testing.T) { + enableAutoscalePrecheckForTest(t) + + params := paramtable.Get() + params.Save(params.QueryCoordCfg.AutoscaleEnabled.Key, "false") + defer params.Reset(params.QueryCoordCfg.AutoscaleEnabled.Key) + + metricsPayload := newAutoscalePrecheckMetrics(t, 4096, 0, 1) + cluster := session.NewMockCluster(t) + cluster.EXPECT(). + GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil). + Once() + server := newAutoscalePrecheckTestServer(t, cluster) + server.broker = newAutoscalePrecheckTestBroker(t, nil, nil, nil) + memoryCounter := metrics.QueryCoordLoadDemandMemoryBytes.WithLabelValues(meta.DefaultResourceGroupName) + diskCounter := metrics.QueryCoordLoadDemandDiskBytes.WithLabelValues(meta.DefaultResourceGroupName) + memoryBefore := testutil.ToFloat64(memoryCounter) + diskBefore := testutil.ToFloat64(diskCounter) + + requiredByRG, err := server.checkLoadResource(context.Background(), newAutoscalePrecheckTestRequest()) + + require.NoError(t, err) + require.Contains(t, requiredByRG, meta.DefaultResourceGroupName) + assert.Equal(t, int64(1024), requiredByRG[meta.DefaultResourceGroupName].memoryBytes) + assert.Zero(t, requiredByRG[meta.DefaultResourceGroupName].diskBytes) + assert.Equal(t, memoryBefore, testutil.ToFloat64(memoryCounter)) + assert.Zero(t, testutil.ToFloat64(diskCounter)-diskBefore) +} + +func TestCheckLoadResourcePreservesResourceInsufficientDecision(t *testing.T) { + enableAutoscalePrecheckForTest(t) + + params := paramtable.Get() + params.Save(params.QueryCoordCfg.AutoscaleEnabled.Key, "false") + defer params.Reset(params.QueryCoordCfg.AutoscaleEnabled.Key) + + metricsPayload := newAutoscalePrecheckMetrics(t, 1024, 1024, 1) + cluster := session.NewMockCluster(t) + cluster.EXPECT(). + GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil). + Once() + server := newAutoscalePrecheckTestServer(t, cluster) + server.broker = newAutoscalePrecheckTestBroker(t, nil, nil, nil) + memoryCounter := metrics.QueryCoordLoadDemandMemoryBytes.WithLabelValues(meta.DefaultResourceGroupName) + diskCounter := metrics.QueryCoordLoadDemandDiskBytes.WithLabelValues(meta.DefaultResourceGroupName) + memoryBefore := testutil.ToFloat64(memoryCounter) + diskBefore := testutil.ToFloat64(diskCounter) + + requiredByRG, err := server.checkLoadResource(context.Background(), newAutoscalePrecheckTestRequest()) + + require.Error(t, err) + assert.Empty(t, requiredByRG) + var precheckErr *loadResourcePrecheckError + require.ErrorAs(t, err, &precheckErr) + status := statusWithLoadResourcePrecheckSuggestion(err) + assert.Equal(t, merr.Code(merr.ErrServiceResourceInsufficient), status.GetCode()) + assert.False(t, status.GetRetriable()) + assert.Contains(t, status.GetExtraInfo(), loadResourcePrecheckSuggestKey) + assert.Equal(t, memoryBefore, testutil.ToFloat64(memoryCounter)) + assert.Equal(t, diskBefore, testutil.ToFloat64(diskCounter)) +} + +func TestQueryNodeResourceSnapshotUsesStaticCapacityAndReturnsMetricsError(t *testing.T) { + params := paramtable.Get() + params.Save(params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key, "0.4") + defer params.Reset(params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key) + params.Save(params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key, "0.5") + defer params.Reset(params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key) + + node := session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1}) + node.UpdateStats(session.WithMemCapacity(1024)) + nodes := []*session.NodeInfo{node} + + server := &Server{} + snapshot, err := server.queryNodeResourceSnapshot(context.Background(), nodes) + require.NoError(t, err) + _, staticCapacity, ok := queryNodeResourceStatusFromSnapshot(nodes, snapshot) + require.True(t, ok) + assert.Positive(t, staticCapacity.memoryBytes) + + cluster := session.NewMockCluster(t) + cluster.EXPECT(). + GetMetrics(mock.Anything, int64(1), mock.Anything). + Return(nil, errors.New("metrics unavailable")). + Once() + server.cluster = cluster + + _, err = server.queryNodeResourceSnapshot(context.Background(), nodes) + require.Error(t, err) +} + +func TestCurrentQueryNodeGlobalCapacityUsesOnlineNodes(t *testing.T) { + params := paramtable.Get() + params.Save(params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key, "0.4") + defer params.Reset(params.QuotaConfig.QueryNodeMemoryLowWaterLevel.Key) + params.Save(params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key, "0.5") + defer params.Reset(params.QuotaConfig.QueryNodeMemoryHighWaterLevel.Key) + params.Save(params.QueryNodeCfg.DiskCapacityLimit.Key, "1") + defer params.Reset(params.QueryNodeCfg.DiskCapacityLimit.Key) + + nodeManager := session.NewNodeManager() + node1 := session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1}) + node1.UpdateStats(session.WithMemCapacity(1024)) + nodeManager.Add(node1) + node2 := session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 2}) + node2.UpdateStats(session.WithMemCapacity(2048)) + nodeManager.Add(node2) + nodeManager.Stopping(2) + node3 := session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 3}) + node3.UpdateStats(session.WithMemCapacity(512)) + nodeManager.Add(node3) + server := &Server{nodeMgr: nodeManager} + + _, _, capacity, err := server.currentQueryNodeResourceStatus(context.Background(), nil, true) + + require.NoError(t, err) + assert.Equal(t, int64((1024+512)*1024*1024), capacity.memoryBytes) + assert.Equal(t, int64(2*bytesPerGB), capacity.diskBytes) +} + +func TestCurrentQueryNodeResourceStatusReusesSnapshotForGlobalCapacity(t *testing.T) { + metricsPayload := newAutoscalePrecheckMetrics(t, 1024, 512, 1) + cluster := session.NewMockCluster(t) + cluster.EXPECT(). + GetMetrics(mock.MatchedBy(func(ctx context.Context) bool { + return retry.MaxAttemptsFromContextOrDefault(ctx, 10) == 1 + }), int64(1), mock.Anything). + Return(&milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil). + Once() + server := newAutoscalePrecheckTestServer(t, cluster) + for i, label := range []string{ + sessionutil.LabelStreamingNodeEmbeddedQueryNode, + sessionutil.LegacyLabelStreamingNodeEmbeddedQueryNode, + } { + server.nodeMgr.Add(session.NewNodeInfo(session.ImmutableNodeInfo{ + NodeID: int64(i + 2), + Labels: map[string]string{label: "1"}, + })) + } + + _, capacityByRG, globalCapacity, err := server.currentQueryNodeResourceStatus( + context.Background(), + map[string]autoscaleResourceUsage{meta.DefaultResourceGroupName: {memoryBytes: 1}}, + true, + ) + + require.NoError(t, err) + assert.Equal(t, autoscaleResourceUsage{memoryBytes: 1024, diskBytes: 1_000_000_000}, globalCapacity) + cluster.AssertNotCalled(t, "GetMetrics", mock.Anything, int64(2), mock.Anything) + cluster.AssertNotCalled(t, "GetMetrics", mock.Anything, int64(3), mock.Anything) + assert.Equal(t, capacityByRG[meta.DefaultResourceGroupName], usableResourceCapacity(globalCapacity)) +} + +func TestQueryNodeResourceSnapshotFetchesMetricsConcurrently(t *testing.T) { + metricsPayload := newAutoscalePrecheckMetrics(t, 1024, 512, 1) + started := make(chan struct{}, 2) + attempts := make(chan uint, 2) + release := make(chan struct{}) + released := false + defer func() { + if !released { + close(release) + } + }() + + cluster := session.NewMockCluster(t) + cluster.EXPECT(). + GetMetrics(mock.Anything, mock.Anything, mock.Anything). + RunAndReturn(func(ctx context.Context, _ int64, _ *milvuspb.GetMetricsRequest) (*milvuspb.GetMetricsResponse, error) { + attempts <- retry.MaxAttemptsFromContextOrDefault(ctx, 10) + started <- struct{}{} + <-release + return &milvuspb.GetMetricsResponse{Status: merr.Success(), Response: metricsPayload}, nil + }). + Twice() + server := &Server{cluster: cluster} + nodes := []*session.NodeInfo{ + session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 1}), + session.NewNodeInfo(session.ImmutableNodeInfo{NodeID: 2}), + } + type snapshotResult struct { + snapshot map[int64]queryNodeResourceStatus + err error + } + resultCh := make(chan snapshotResult, 1) + go func() { + snapshot, err := server.queryNodeResourceSnapshot(context.Background(), nodes) + resultCh <- snapshotResult{snapshot: snapshot, err: err} + }() + + for range 2 { + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("query node metrics requests did not run concurrently") + } + } + close(release) + released = true + for range 2 { + assert.Equal(t, uint(1), <-attempts) + } + select { + case result := <-resultCh: + require.NoError(t, result.err) + assert.Len(t, result.snapshot, 2) + case <-time.After(time.Second): + t.Fatal("query node metrics snapshot did not complete") + } +} + +func TestAutoscalePrecheckStatusExtraInfo(t *testing.T) { + err := &loadResourcePrecheckError{ + message: "Insufficient cluster resources. Please scale out the cluster or enable autoscaling. Suggested scale-out ratio: 75%", + suggestedExpandPercent: 75, + } + + status := statusWithLoadResourcePrecheckSuggestion(err) + + assert.Equal(t, merr.Code(merr.ErrServiceResourceInsufficient), status.GetCode()) + assert.Equal(t, "Insufficient cluster resources. Please scale out the cluster or enable autoscaling. Suggested scale-out ratio: 75%", status.GetReason()) + require.NotNil(t, status.GetExtraInfo()) + assert.Equal(t, "75", status.GetExtraInfo()["suggested_expand_percent"]) + assert.NotContains(t, status.GetExtraInfo(), "memory_bytes") + assert.NotContains(t, status.GetExtraInfo(), "disk_bytes") +} diff --git a/internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go b/internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go index cc968d71c1..4fb5055c96 100644 --- a/internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go +++ b/internal/querycoordv2/ddl_callbacks_alter_load_info_load_collection.go @@ -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 diff --git a/internal/querycoordv2/ddl_callbacks_alter_load_info_load_partitions.go b/internal/querycoordv2/ddl_callbacks_alter_load_info_load_partitions.go index 7e766f4920..254881b49a 100644 --- a/internal/querycoordv2/ddl_callbacks_alter_load_info_load_partitions.go +++ b/internal/querycoordv2/ddl_callbacks_alter_load_info_load_partitions.go @@ -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 } diff --git a/internal/querycoordv2/ddl_callbacks_alter_load_info_test.go b/internal/querycoordv2/ddl_callbacks_alter_load_info_test.go index 8aca54b657..0319ffac84 100644 --- a/internal/querycoordv2/ddl_callbacks_alter_load_info_test.go +++ b/internal/querycoordv2/ddl_callbacks_alter_load_info_test.go @@ -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(). diff --git a/internal/querycoordv2/meta/resource_manager.go b/internal/querycoordv2/meta/resource_manager.go index 1941af3e97..d19c14687c 100644 --- a/internal/querycoordv2/meta/resource_manager.go +++ b/internal/querycoordv2/meta/resource_manager.go @@ -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. diff --git a/internal/querycoordv2/meta/resource_manager_test.go b/internal/querycoordv2/meta/resource_manager_test.go index 1a920f07d3..43b60f0995 100644 --- a/internal/querycoordv2/meta/resource_manager_test.go +++ b/internal/querycoordv2/meta/resource_manager_test.go @@ -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)) diff --git a/internal/querycoordv2/server_test.go b/internal/querycoordv2/server_test.go index ac2343093b..489dfcb1a1 100644 --- a/internal/querycoordv2/server_test.go +++ b/internal/querycoordv2/server_test.go @@ -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) { diff --git a/internal/querycoordv2/services.go b/internal/querycoordv2/services.go index ad2a2466be..19be2d674e 100644 --- a/internal/querycoordv2/services.go +++ b/internal/querycoordv2/services.go @@ -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() diff --git a/internal/querycoordv2/services_test.go b/internal/querycoordv2/services_test.go index 3ba84ae8e6..afc8c4fde0 100644 --- a/internal/querycoordv2/services_test.go +++ b/internal/querycoordv2/services_test.go @@ -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() { diff --git a/internal/querycoordv2/task/utils.go b/internal/querycoordv2/task/utils.go index f9a7f0ceed..e7d1e76dfd 100644 --- a/internal/querycoordv2/task/utils.go +++ b/internal/querycoordv2/task/utils.go @@ -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 diff --git a/internal/querynodev2/metrics_info.go b/internal/querynodev2/metrics_info.go index 477274a72d..c1e4b56349 100644 --- a/internal/querynodev2/metrics_info.go +++ b/internal/querynodev2/metrics_info.go @@ -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) diff --git a/internal/querynodev2/metrics_info_test.go b/internal/querynodev2/metrics_info_test.go index 98d2bd41bf..55f701f167 100644 --- a/internal/querynodev2/metrics_info_test.go +++ b/internal/querynodev2/metrics_info_test.go @@ -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)) +} diff --git a/internal/querynodev2/segments/load_index_info.go b/internal/querynodev2/segments/load_index_info.go deleted file mode 100644 index dc241cbd8c..0000000000 --- a/internal/querynodev2/segments/load_index_info.go +++ /dev/null @@ -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 -#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) -} diff --git a/internal/querynodev2/segments/mock_loader.go b/internal/querynodev2/segments/mock_loader.go index d07473295a..c3d346a5f3 100644 --- a/internal/querynodev2/segments/mock_loader.go +++ b/internal/querynodev2/segments/mock_loader.go @@ -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)) diff --git a/internal/querynodev2/segments/segment.go b/internal/querynodev2/segments/segment.go index c000ffae32..6b885284fd 100644 --- a/internal/querynodev2/segments/segment.go +++ b/internal/querynodev2/segments/segment.go @@ -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() { diff --git a/internal/querynodev2/segments/segment_loader.go b/internal/querynodev2/segments/segment_loader.go index a084b356bc..493549595f 100644 --- a/internal/querynodev2/segments/segment_loader.go +++ b/internal/querynodev2/segments/segment_loader.go @@ -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 diff --git a/internal/querynodev2/segments/segment_loader_test.go b/internal/querynodev2/segments/segment_loader_test.go index f83ffaebf6..4a0c731a8b 100644 --- a/internal/querynodev2/segments/segment_loader_test.go +++ b/internal/querynodev2/segments/segment_loader_test.go @@ -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()) diff --git a/internal/querynodev2/segments/segment_test.go b/internal/querynodev2/segments/segment_test.go index 8fd3128aa6..29d3400c02 100644 --- a/internal/querynodev2/segments/segment_test.go +++ b/internal/querynodev2/segments/segment_test.go @@ -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() { diff --git a/internal/querynodev2/segments/utils.go b/internal/querynodev2/segments/utils.go index ad44805474..f75c0ba663 100644 --- a/internal/querynodev2/segments/utils.go +++ b/internal/querynodev2/segments/utils.go @@ -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 { diff --git a/internal/querynodev2/segments/utils_test.go b/internal/querynodev2/segments/utils_test.go index 2b6c817db0..ecd3eb6d95 100644 --- a/internal/querynodev2/segments/utils_test.go +++ b/internal/querynodev2/segments/utils_test.go @@ -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() diff --git a/internal/util/segcore/loadresource/index.go b/internal/util/segcore/loadresource/index.go new file mode 100644 index 0000000000..9b2e6409ff --- /dev/null +++ b/internal/util/segcore/loadresource/index.go @@ -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 +#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), + } +} diff --git a/internal/util/segcore/loadresource/index_test.go b/internal/util/segcore/loadresource/index_test.go new file mode 100644 index 0000000000..83abc84214 --- /dev/null +++ b/internal/util/segcore/loadresource/index_test.go @@ -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) +} diff --git a/internal/util/segcore/loadresource/segment.go b/internal/util/segcore/loadresource/segment.go new file mode 100644 index 0000000000..f35db0dbe7 --- /dev/null +++ b/internal/util/segcore/loadresource/segment.go @@ -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 := ¶mtable.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 := ¶mtable.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)) +} diff --git a/internal/util/segcore/loadresource/segment_test.go b/internal/util/segcore/loadresource/segment_test.go new file mode 100644 index 0000000000..7d14bafbbc --- /dev/null +++ b/internal/util/segcore/loadresource/segment_test.go @@ -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()) +} diff --git a/pkg/metrics/metrics_test.go b/pkg/metrics/metrics_test.go index 40c5701c0a..5b4039a29f 100644 --- a/pkg/metrics/metrics_test.go +++ b/pkg/metrics/metrics_test.go @@ -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) diff --git a/pkg/metrics/querycoord_metrics.go b/pkg/metrics/querycoord_metrics.go index 20c30b820b..be9fbedbc2 100644 --- a/pkg/metrics/querycoord_metrics.go +++ b/pkg/metrics/querycoord_metrics.go @@ -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) diff --git a/pkg/util/metricsinfo/metrics_info.go b/pkg/util/metricsinfo/metrics_info.go index 730287d3f6..720193ac0e 100644 --- a/pkg/util/metricsinfo/metrics_info.go +++ b/pkg/util/metricsinfo/metrics_info.go @@ -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. diff --git a/pkg/util/metricsinfo/metrics_info_test.go b/pkg/util/metricsinfo/metrics_info_test.go index 0f51df4233..8079494dd7 100644 --- a/pkg/util/metricsinfo/metrics_info_test.go +++ b/pkg/util/metricsinfo/metrics_info_test.go @@ -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) diff --git a/pkg/util/paramtable/component_param.go b/pkg/util/paramtable/component_param.go index 1b79aea9e6..280e710119 100644 --- a/pkg/util/paramtable/component_param.go +++ b/pkg/util/paramtable/component_param.go @@ -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", diff --git a/pkg/util/paramtable/component_param_test.go b/pkg/util/paramtable/component_param_test.go index 22a447f32d..1269403dd7 100644 --- a/pkg/util/paramtable/component_param_test.go +++ b/pkg/util/paramtable/component_param_test.go @@ -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) {